Leet Code Daily 2024-06-26 - Balance a Binary Search Tree

/**
 * 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:
    void collect(TreeNode* root, vector<TreeNode*>& out) {
        if (root == nullptr) {
            return;
        }
        collect(root->left, out);
        out.push_back(root);
        collect(root->right, out);
    }

    TreeNode* rebuild(vector<TreeNode*>& nodes, int l, int r) {
        if (l > r) {
            return nullptr;
        }
        int mid = (l + r) >> 1;
        nodes[mid]->left = rebuild(nodes, l, mid - 1);
        nodes[mid]->right = rebuild(nodes, mid + 1, r);
        return nodes[mid];
    }

    TreeNode* balanceBST(TreeNode* root) {
        vector<TreeNode*> nodes;
        collect(root, nodes);
        return rebuild(nodes, 0, nodes.size() - 1);
    }
};