Mind palace

some thoughts

0%

1382. Balance a Binary Search Tree

Posted on Edited on In Disqus:
Symbols count in article: 2.1k Reading time ≈ 2 mins.

O(n) time O(n) space
先dfs扫一遍得到所有的node然后二分构造bst

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    TreeNode* balanceBST(TreeNode* root) {
        dfs(root);
        return reconstruct(0, nodes.size());
    }

    void dfs(TreeNode *root) {
        if (!root) return;
        dfs(root->left);
        root->left = nullptr;
        nodes.push_back(root);
        auto r = root->right;
        root->right = nullptr;
        dfs(r);
    }

    TreeNode *reconstruct(int b, int e) {
        if (b == e) return nullptr;
        int m = (b + e) / 2;
        auto res = nodes[m];
        res->left = reconstruct(b, m);
        res->right = reconstruct(m + 1, e);
        return res;
    }

    vector<TreeNode *> nodes;
};

DSW O(n) time O(1) space
解法
先通过多次右旋变成right-skew的单链，再按照树高（以及子树高）多次左旋

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    int makeVine(TreeNode *grand, int cnt = 0) {
      auto n = grand->right;
      while (n != nullptr) {
        if (n->left != nullptr) {
          auto old_n = n;
          n = n->left;
          old_n->left = n->right;
          n->right = old_n;
          grand->right = n;
        }
        else {
            ++cnt;
            grand = n;
            n = n->right;
        }
      }
      return cnt;
    }
    void compress(TreeNode *grand, int m) {
      auto n = grand->right;
      while (m-- > 0) {
        auto old_n = n;
        n = n->right;
        grand->right = n;
        old_n->right = n->left;
        n->left = old_n;
        grand = n;
        n = n->right;
      }
    }
    TreeNode* balanceBST(TreeNode *root) {
      TreeNode grand;
      grand.right = root;
      auto cnt = makeVine(&grand);
      int m = pow(2, int(log2(cnt + 1))) - 1;
      compress(&grand, cnt - m);
      for (m = m / 2; m > 0; m /= 2)
        compress(&grand, m);
      return grand.right;
    }
};