919. Complete Binary Tree Inserter

https://leetcode.com/problems/complete-binary-tree-inserter/description/

A complete binary tree is a binary tree in which every level, except possibly the last, is completely filled, and all nodes are as far left as possible.

Write a data structure CBTInserter that is initialized with a complete binary tree and supports the following operations:

• CBTInserter(TreeNode root) initializes the data structure on a given tree with head node root;

• CBTInserter.insert(int v) will insert a TreeNode into the tree with value node.val = v so that the tree remains complete, and returns the value of the parent of the inserted TreeNode;

• CBTInserter.get_root() will return the head node of the tree.

Example 1:

Input: inputs = ["CBTInserter","insert","get_root"], inputs = [[[1]],[2],[]]
Output: [null,1,[1,2]]

Example 2:

Input: inputs = ["CBTInserter","insert","insert","get_root"], inputs = [[[1,2,3,4,5,6]],[7],[8],[]]
Output: [null,3,4,[1,2,3,4,5,6,7,8]]

Thoughts

实现初始和插入完全二叉树。完全二叉树存储方法是数组。2i + 1是左儿子，2i + 2是右儿子。因此(I - 1) / 2是parent，当插入的新node i%2是1时为parent的左节点，否则是右节点。

Code

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class CBTInserter:

    def __init__(self, root: TreeNode):
        self.a = [root]
        for node in self.a:
            if node.left: self.a.append(node.left)
            if node.right: self.a.append(node.right)
        
    def insert(self, v: int) -> int:
        a, N = self.a, len(self.a)
        a.append(TreeNode(v))
        p = a[(N - 1) // 2]
        if N % 2: p.left = a[-1]
        else: p.right = a[-1]
        return p.val

    def get_root(self) -> TreeNode:
        return self.a[0]

# Your CBTInserter object will be instantiated and called as such:
# obj = CBTInserter(root)
# param_1 = obj.insert(v)
# param_2 = obj.get_root()

class CBTInserter {
public:
    CBTInserter(TreeNode* root) {
        tree.push_back(root);
        for (int i = 0; i < tree.size(); ++i) {
            if (tree[i]->left) tree.push_back(tree[i]->left);
            if (tree[i]->right) tree.push_back(tree[i]->right);
        }
    }

    int insert(int v) {
        TreeNode *node = new TreeNode(v);
        const auto i = tree.size();
        const auto p = tree[(i - 1) / 2];
        if (i % 2) {
            p->left = node;
        } else {
            p->right = node;
        }
        tree.push_back(node);
        return p->val;
    }

    TreeNode* get_root() {
        return tree[0];
    }

private: 
    vector<TreeNode*> tree;
};