BST — 進階模式與深入探討

# python
# LC 173 - Binary Search Tree Iterator
# IDEA: keep the unvisited left spine on a stack; pop = next value, then push the right child's spine
# time = O(1) amortised per next(), space = O(h)
class BSTIterator:
    def __init__(self, root):
        self.stack = []
        self._push_left(root)

    def _push_left(self, node):
        while node:                       # the smallest unvisited value ends up on top
            self.stack.append(node)
            node = node.left

    def next(self):
        node = self.stack.pop()
        self._push_left(node.right)       # its right subtree comes next in order
        return node.val

    def hasNext(self):
        return len(self.stack) > 0

// java
// LC 173 - Binary Search Tree Iterator
// IDEA: unvisited left spine on a stack; pop = next value, then push the right child's spine
// time = O(1) amortised per next(), space = O(h)
class BSTIterator {
    private final Deque<TreeNode> stack = new ArrayDeque<>();

    public BSTIterator(TreeNode root) {
        pushLeft(root);
    }

    private void pushLeft(TreeNode node) {
        while (node != null) {
            stack.push(node);
            node = node.left;
        }
    }

    public int next() {
        TreeNode node = stack.pop();
        pushLeft(node.right);
        return node.val;
    }

    public boolean hasNext() {
        return !stack.isEmpty();
    }
}

init        push 7, 3                stack [7, 3]
next() = 3  pop 3, no right child    stack [7]
next() = 7  pop 7, push 15, 9        stack [15, 9]
next() = 9  pop 9                    stack [15]
next() = 15 pop 15, push 20          stack [20]
next() = 20 pop 20                   stack []      hasNext() = false

# python
# LC 1305 - All Elements in Two Binary Search Trees
# IDEA: two left-spine stacks; pop from the one whose top is smaller, then push its right spine
# time = O(m + n), space = O(h1 + h2) excluding the output
def getAllElements(root1, root2):
    stack1, stack2, result = [], [], []

    def push_left(stack, node):
        while node:
            stack.append(node)
            node = node.left

    push_left(stack1, root1)
    push_left(stack2, root2)
    while stack1 or stack2:
        # take from stack1 when stack2 is empty, or when both are live and stack1's top is smaller
        if not stack2 or (stack1 and stack1[-1].val <= stack2[-1].val):
            node = stack1.pop()
            push_left(stack1, node.right)
        else:
            node = stack2.pop()
            push_left(stack2, node.right)
        result.append(node.val)
    return result

# python
# LC 99 - Recover Binary Search Tree
# IDEA: in-order walk with a prev pointer; first = prev at the first drop, second = node at the last
# time = O(n), space = O(h)
def recoverTree(root):
    first = second = prev = None

    def inorder(node):
        nonlocal first, second, prev
        if not node:
            return
        inorder(node.left)
        if prev and prev.val > node.val:     # a drop in the sorted order
            if first is None:
                first = prev                 # the too-large value, from the first drop only
            second = node                    # the too-small value, from the last drop
        prev = node
        inorder(node.right)

    inorder(root)
    first.val, second.val = second.val, first.val   # swap values; the structure stays

// java
// LC 99 - Recover Binary Search Tree
// IDEA: in-order with prev; first is set once, second at every drop; then swap the values
// time = O(n), space = O(h)
private TreeNode first, second, prev;

public void recoverTree(TreeNode root) {
    first = second = prev = null;
    inorder(root);
    int temp = first.val;
    first.val = second.val;
    second.val = temp;
}

private void inorder(TreeNode node) {
    if (node == null) {
        return;
    }
    inorder(node.left);
    if (prev != null && prev.val > node.val) {
        if (first == null) {
            first = prev;
        }
        second = node;
    }
    prev = node;
    inorder(node.right);
}

# python
# LC 99 - Recover Binary Search Tree, O(1) extra space
# IDEA: Morris in-order (thread the predecessor, unthread on the second visit) + the same drop rule
# time = O(n), space = O(1)
def recoverTree(root):
    first = second = prev = None
    cur = root
    while cur:
        if cur.left:
            pred = cur.left
            while pred.right and pred.right is not cur:
                pred = pred.right
            if not pred.right:               # first visit: thread, then go left
                pred.right = cur
                cur = cur.left
                continue
            pred.right = None                # second visit: unthread, then visit cur
        if prev and prev.val > cur.val:
            if first is None:
                first = prev
            second = cur
        prev = cur
        cur = cur.right
    first.val, second.val = second.val, first.val

6 > 3   first drop  -> first = 6, second = 3
5 > 2   second drop -> first stays 6, second = 2
swap 6 and 2        -> [1, 2, 3, 4, 5, 6, 7]

# python
# LC 703 - Kth Largest Element in a Stream
# IDEA: each node keeps its subtree size; kth largest descends using size(node.right)
# time = O(h) per add, space = O(n)
class CountNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None
        self.count = 1                       # nodes in the subtree rooted here


class KthLargest:
    def __init__(self, k, nums):
        self.k = k
        self.root = None
        for v in nums:
            self.root = self._insert(self.root, v)

    def _insert(self, node, val):
        if not node:
            return CountNode(val)
        node.count += 1                      # the new value lands somewhere below
        if val < node.val:
            node.left = self._insert(node.left, val)
        else:
            node.right = self._insert(node.right, val)   # duplicates go right
        return node

    def _size(self, node):
        return node.count if node else 0

    def add(self, val):
        self.root = self._insert(self.root, val)
        node, k = self.root, self.k
        while True:
            right_size = self._size(node.right)          # values larger than node.val
            if k <= right_size:
                node = node.right
            elif k == right_size + 1:
                return node.val
            else:
                k -= right_size + 1                      # skip the right subtree and node
                node = node.left

// java
// LC 703 - Kth Largest Element in a Stream
// IDEA: subtree sizes on every node; rank-descend using the right subtree's size
// time = O(h) per add, space = O(n)
class KthLargest {
    private static class Node {
        int val;
        int count = 1;
        Node left, right;
        Node(int val) { this.val = val; }
    }

    private final int k;
    private Node root;

    public KthLargest(int k, int[] nums) {
        this.k = k;
        for (int v : nums) {
            root = insert(root, v);
        }
    }

    private Node insert(Node node, int val) {
        if (node == null) {
            return new Node(val);
        }
        node.count++;
        if (val < node.val) {
            node.left = insert(node.left, val);
        } else {
            node.right = insert(node.right, val);
        }
        return node;
    }

    private int size(Node node) {
        return node == null ? 0 : node.count;
    }

    public int add(int val) {
        root = insert(root, val);
        Node node = root;
        int rank = k;
        while (true) {
            int rightSize = size(node.right);
            if (rank <= rightSize) {
                node = node.right;
            } else if (rank == rightSize + 1) {
                return node.val;
            } else {
                rank -= rightSize + 1;
                node = node.left;
            }
        }
    }
}

tree (count):  4(5) -> left 2(2) -> right 3(1);  right 5(2) -> right 8(1)
at 4: size(right) = 2, k = 3 == 2 + 1  -> answer 4

# python
# LC 1008 - Construct Binary Search Tree from Preorder Traversal
# IDEA: one cursor over the preorder; a call builds a node only if the next value fits (low, high)
# time = O(n), space = O(h)
def bstFromPreorder(preorder):
    index = 0

    def build(low, high):
        nonlocal index
        if index == len(preorder) or not (low < preorder[index] < high):
            return None                          # the next value belongs to an ancestor
        node = TreeNode(preorder[index])
        index += 1
        node.left = build(low, node.val)         # smaller values, until one is too big
        node.right = build(node.val, high)
        return node

    return build(float("-inf"), float("inf"))

// java
// LC 1008 - Construct Binary Search Tree from Preorder Traversal
// IDEA: a shared cursor; build a node only while the next value fits the interval
// time = O(n), space = O(h)
private int index;

public TreeNode bstFromPreorder(int[] preorder) {
    index = 0;
    return build(preorder, Long.MIN_VALUE, Long.MAX_VALUE);
}

private TreeNode build(int[] preorder, long low, long high) {
    if (index == preorder.length || preorder[index] <= low || preorder[index] >= high) {
        return null;
    }
    TreeNode node = new TreeNode(preorder[index++]);
    node.left = build(preorder, low, node.val);
    node.right = build(preorder, node.val, high);
    return node;
}

build(-inf, inf)   takes 8
  build(-inf, 8)   takes 5
    build(-inf, 5) takes 1;  its children see 7, which is outside (-inf, 1) and (1, 5) -> None
    build(5, 8)    takes 7;  its children see 10, outside (5, 7) and (7, 8)       -> None
  build(8, inf)    takes 10; left sees 12, outside (8, 10) -> None; right takes 12
result [8,5,10,1,7,null,12]

# python
# LC 109 - Convert Sorted List to Binary Search Tree
# IDEA: slow/fast pointers find the middle of [left, right); it is the root, recurse on both halves
# time = O(n log n), space = O(log n)
def sortedListToBST(head):
    def middle(left, right):
        slow = fast = left
        while fast is not right and fast.next is not right:
            slow = slow.next
            fast = fast.next.next
        return slow

    def convert(left, right):                  # builds the nodes in [left, right)
        if left is right:
            return None
        mid = middle(left, right)
        node = TreeNode(mid.val)
        node.left = convert(left, mid)
        node.right = convert(mid.next, right)
        return node

    return convert(head, None)

# python
# LC 109 - in-order simulation
# IDEA: build(lo, hi) makes the left subtree first, then takes the list's current head as its root
# time = O(n), space = O(log n)
def sortedListToBST(head):
    n = 0
    node = head
    while node:
        n += 1
        node = node.next
    current = head

    def build(lo, hi):
        nonlocal current
        if lo > hi:
            return None
        mid = (lo + hi) // 2
        left = build(lo, mid - 1)              # consumes exactly the values before mid
        root = TreeNode(current.val)
        current = current.next
        root.left = left
        root.right = build(mid + 1, hi)
        return root

    return build(0, n - 1)

# python
# LC 1382 - Balance a Binary Search Tree
# IDEA: in-order into a sorted list, then rebuild with the middle of each range as its root
# time = O(n), space = O(n)
def balanceBST(root):
    values = []

    def inorder_walk(node):
        if not node:
            return
        inorder_walk(node.left)
        values.append(node.val)
        inorder_walk(node.right)

    def build(lo, hi):
        if lo > hi:
            return None
        mid = (lo + hi) // 2
        node = TreeNode(values[mid])
        node.left = build(lo, mid - 1)
        node.right = build(mid + 1, hi)
        return node

    inorder_walk(root)
    return build(0, len(values) - 1)

# python
# LC 95 - Unique Binary Search Trees II
# IDEA: each value in [lo, hi] as the root; pair every left subtree with every right subtree
# time = O(n * Catalan(n)), space = O(n * Catalan(n)) for the output
def generateTrees(n):
    def generate(lo, hi):
        if lo > hi:
            return [None]                        # one empty tree, so the pairing still runs
        trees = []
        for root_val in range(lo, hi + 1):
            for left in generate(lo, root_val - 1):
                for right in generate(root_val + 1, hi):
                    trees.append(TreeNode(root_val, left, right))
        return trees

    return generate(1, n) if n else []

// java
// LC 95 - Unique Binary Search Trees II
// IDEA: each value as the root; the cartesian product of the left and right subtree lists
// time = O(n * Catalan(n)), space = O(n * Catalan(n))
public List<TreeNode> generateTrees(int n) {
    return n == 0 ? new ArrayList<>() : generate(1, n);
}

private List<TreeNode> generate(int lo, int hi) {
    List<TreeNode> trees = new ArrayList<>();
    if (lo > hi) {
        trees.add(null);
        return trees;
    }
    for (int rootVal = lo; rootVal <= hi; rootVal++) {
        for (TreeNode left : generate(lo, rootVal - 1)) {
            for (TreeNode right : generate(rootVal + 1, hi)) {
                trees.add(new TreeNode(rootVal, left, right));
            }
        }
    }
    return trees;
}

# python
# LC 96 - Unique Binary Search Trees
# IDEA: dp[nodes] = sum over every root of dp[left size] * dp[right size]
# time = O(n^2), space = O(n)
def numTrees(n):
    dp = [0] * (n + 1)
    dp[0] = 1                                    # the empty tree
    for nodes in range(1, n + 1):
        for root in range(1, nodes + 1):
            dp[nodes] += dp[root - 1] * dp[nodes - root]
    return dp[n]