1382-kir3i.md

1382. Balance a Binary Search Tree

Solution

• 시간복잡도: O(NlogN)

• 알고리즘

  자료구조, 구현

• 풀이설명

  1. 주어진 BST를 중위순회하면 정렬된 상태의 리스트(vals)로 만들 수 있습니다.
  2. BST를 만들 때 중위값부터 넣으면 balanced BST를 만들 수 있으므로, 재귀함수로 새로운 BST에 계속 중위값을 넣어줍니다.

• 소스코드

  • C++

    class Solution {
    public:
        vector<int> vals;
        TreeNode* balanceBST(TreeNode* root) {
            popNode(root);
            return makeTree(0, vals.size() - 1);
        }
        
        TreeNode* makeTree(int lo, int hi) {
            if (lo > hi)        return nullptr;
            else if (lo == hi)  return new TreeNode(vals[lo]);
            
            int mid = (lo + hi) / 2;
            TreeNode* root = new TreeNode(vals[mid]);
            root->left = makeTree(lo, mid - 1);
            root->right = makeTree(mid + 1, hi);
            return root;
        }
        
        void popNode(TreeNode *root) {
            if(!root)   return;
            popNode(root->left);
            vals.push_back(root->val);
            popNode(root->right);
        }
    };

  • Python

    class Solution:
        def balanceBST(self, root):
            vals = []
            self.popNode(root, vals)
            return self.makeTree(0, len(vals)-1, vals)
    
        def makeTree(self, lo, hi, vals):
            if lo > hi:
                return None
            elif lo == hi:
                return TreeNode(vals[lo])
            mid = (lo + hi) // 2
            root = TreeNode(vals[mid])
            root.left = self.makeTree(lo, mid-1, vals)
            root.right = self.makeTree(mid+1, hi, vals)
    
            return root
    
        def popNode(self, root, res):
            if root.left:
                self.popNode(root.left, res)
            res.append(root.val)
            if root.right:
                self.popNode(root.right, res)