leetcode

0 class Solution {

1 struct record {

2 TreeNode *root;

3 int low, high;

4 record(TreeNode *root, int low, int high) : root(root), low(low), high(high) {}

5 };

6

7 public:

8 TreeNode *balanceBST(TreeNode *root) {

9 vector<TreeNode *> nums;

10

11 {

12 stack<TreeNode *> st;

13 while (true) {

14 while (root) {

15 st.push(root);

16 root = root->left;

17 }

18 if (st.empty()) break;

19 root = st.top(), st.pop();

20 nums.push_back(root);

21 root = root->right;

22 }

23 }

24

25 stack<record> st;

26 TreeNode *head = new TreeNode(INT_MIN), *t;

27 st.push({head, 0, (int)nums.size() - 1});

28 while (!st.empty()) {

29 record r = st.top();

30 st.pop();

31 while (r.low <= r.high) {

32 int mid = r.low + (r.high - r.low) / 2;

33 nums[mid]->left = nums[mid]->right = nullptr;

34 (nums[mid]->val >= r.root->val ? r.root->right : r.root->left) = t = nums[mid];

35 st.push({r.root = t, mid + 1, r.high});

36 r.high = mid - 1;

37 }

38 }

39 return head->right;

40 }

41 };