findClosestLeaf.ts

/**
 * Definition for a binary tree node.
 * class TreeNode {
 *     val: number
 *     left: TreeNode | null
 *     right: TreeNode | null
 *     constructor(val?: number, left?: TreeNode | null, right?: TreeNode | null) {
 *         this.val = (val===undefined ? 0 : val)
 *         this.left = (left===undefined ? null : left)
 *         this.right = (right===undefined ? null : right)
 *     }
 * }
 */

function findClosestLeaf(root: TreeNode | null, k: number): number {
    if (!root.left && !root.right) {
        return k;
    }

    // BFS or DFS doesn't matter here because we need to crawl entire graph
    // I reckon DFS is faster in JS with Array.prototype.pop()
    const stack: TreeNode[] = [];
    // Easy access to nodes in new graph
    const graphMap = new Map<number, GraphNode>();
    const leafNodes: GraphNode[] = [];

    stack.push(root);
    graphMap.set(root.val, {val: root.val, parent: null, level: 0});

    // Generate the graph of nodes with parents + level
    while (stack.length) {
        const { val: currVal, left, right } = stack.pop();
        const isLeaf = !left && !right;
        const graphNode = graphMap.get(currVal);

        if (isLeaf) {
            // return early if K is a leaf
            if (currVal === k) {
                return k;
            }

            // keep track of leaf nodes
            leafNodes.push(graphNode);
        }

        if (left) {
            const leftGraphNode = {val: left.val, parent: graphNode, level: graphNode.level + 1};
            graphMap.set(left.val, leftGraphNode);
            stack.push(left);
        }

        if (right) {
            const rightGraphNode = {val: right.val, parent: graphNode, level: graphNode.level + 1};
            graphMap.set(right.val, rightGraphNode);
            stack.push(right);
        }
    }

    // Keep track of all ancesotrs of k
    const ancestors = new Map<number, GraphNode>();
    const kNode = graphMap.get(k);
    let node = kNode;

    // Hydrate the ancestors map
    while (node) {
        ancestors.set(node.val, node);
        node = node.parent;
    }

    // Keep track of current best leaf and best distance
    let bestLeaf: GraphNode | null;
    let bestDistance = Infinity;

    // We are going to BFS each leaf to a common ancestor because DFS might be painful
    // in the case of [1,2,3,4,null,null,5,null,null,null,6,null,7] k=2
    const queue = new BfsQueue<{ leaf: GraphNode, ancestor: GraphNode }>();
    // Queue is of leaf node + current ancestor we BFS-ing up
    leafNodes.forEach(leaf => queue.enqueue({leaf, ancestor: leaf.parent}));

    // BFS for an ancestor that is in k's ancestor chain
    while (queue.size) {
        const { leaf, ancestor } = queue.dequeue();

        if (ancestors.has(ancestor.val)) {
            // calculate the distance between the nodes
            // AKA (kNode.level - ancestor.level) + (leaf.level - ancestor.level)
            const distance = kNode.level + leaf.level - 2 * ancestor.level;
            if (distance < bestDistance) {
                bestLeaf = leaf;
                bestDistance = distance;
            }

            continue;
        }

        queue.enqueue({leaf, ancestor: ancestor.parent});
    }

    return bestLeaf.val;;
};

// standard queue implementation for BFS. Can't call it Queue because it explodes leetcode tisk tisk
class BfsQueue<T> {
    arr: {[idx: number]: T} = {};
    private startIndex = 0;
    private endIndex = 0;

    get size() {
        return this.endIndex - this.startIndex;
    }

    enqueue(item: T) {
        this.arr[this.endIndex] = item;
        this.endIndex++;
    }

    dequeue() {
        const val = this.arr[this.startIndex];
        delete this.arr[this.startIndex];
        this.startIndex++;
        return val;
    }
}

interface GraphNode {
    val: number;
    parent: GraphNode | null;
    level: number;
}