3319.k-th-largest-perfect-subtree-size-in-binary-tree.md

3319. K-th Largest Perfect Subtree Size in Binary Tree

Bottom-Up DFS + Heap

We can use the bottom-up DFS similar to 110. Balanced Binary Tree to calculate the depth of each node:

• We calculate the depth of subtree recursively.
• If the subtree is a perfect subtree, then we have to ensure the current root is also a perfect tree.
  • The depth of all the subtrees should be the same.
  • The number of nodes in the subtree should be 2^n - 1, where n is the depth.
• Otherwise, if there is a subtree that is not a perfect subtree, we return -1 to indicate the node is not a perfect subtree.

     node          node
   /      \      /      \
  2        1    -1       1

Then we use a fixed-size k min heap to store the size of the perfect subtree so that we can get the kth largest perfect subtree size.

We can also use number of nodes to determine if the subtree is a perfect subtree.

fun kthLargestPerfectSubtree(root: TreeNode?, k: Int): Int {
    val minHeap = PriorityQueue<Int>()
    dfs(root, minHeap, k)
    return if (minHeap.size == k) minHeap.peek() else -1
}

// Return the depth
private fun dfs(root: TreeNode?, minHeap: PriorityQueue<Int>, k: Int): Int {
    if (root == null) return 0
    val left = dfs(root.left, minHeap, k)
    val right = dfs(root.right, minHeap, k)

    if (left != -1 && right != -1) {
        if (left == right) {
            // number of nodes are 2^n - 1, where `n` is the depth
            val depth = left + 1
            val nodesCount = Math.pow(2.0, depth.toDouble()).toInt() - 1
            minHeap.add(nodesCount, k)
            return depth
        }
    }
    return -1
}

private fun PriorityQueue<Int>.add(count: Int, k: Int) {
    this.add(count)
    if (this.size > k) {
        this.poll()
    }
}