Problem Description

You are given the root of a binary tree and an integer k. Your task is to count how many nodes in the tree are "great enough."

A node is considered great enough if it satisfies both of these conditions:

1. Its subtree (including the node itself) contains at least k nodes
2. The node's value is greater than at least k nodes' values in its subtree

The subtree of a node includes the node itself and all its descendants.

For example, if k = 3, a node is great enough if:

• Its subtree has 3 or more nodes (counting itself and all descendants)
• Its value is larger than at least 3 nodes' values in that subtree

The solution uses a depth-first search (DFS) approach with a min-heap to efficiently track the k largest values in each subtree. The heap maintains at most k elements, where the smallest of these k largest values is at the top. By negating values when storing them, the code converts Python's min-heap into effectively tracking the largest values. When a node's value is greater than the smallest of the k largest values in its subtree (checked via -l[0] < root.val), and the subtree has exactly k nodes, the node is counted as great enough.

The function returns the total count of nodes that meet both criteria.

Flowchart Walkthrough

First, let's pin down the algorithm using the Flowchart. Here's a step-by-step walkthrough:

Is it a graph?

• Yes: While this is a binary tree problem, a tree is a special type of graph (specifically, a connected acyclic graph). Binary trees are graphs with nodes and edges where each node has at most two children.

Is it a tree?

• Yes: The problem explicitly states we're working with a binary tree. We need to traverse the tree and analyze properties of each node and its subtree.

DFS

• Yes: We arrive at DFS (Depth-First Search) as our algorithm choice.

Conclusion: The flowchart suggests using a DFS approach for this problem.

This makes perfect sense because:

1. We need to process each node's entire subtree before determining if that node is "great enough"
2. We need to collect information from children before processing the parent (post-order traversal)
3. DFS allows us to recursively compute subtree properties - gathering the count of nodes and their values from each subtree
4. The bottom-up nature of DFS lets us build up the k largest values from leaves to root, which is exactly what we need to check if a node's value is greater than at least k nodes in its subtree

The solution implements DFS by recursively traversing left and right subtrees, maintaining a heap of the k largest values seen so far, and counting nodes that satisfy both conditions (having at least k nodes in subtree and value greater than at least k nodes).

Intuition

When we need to determine if a node is "great enough," we must know two things about its subtree: how many nodes it contains and what values those nodes have. This naturally suggests a bottom-up approach where we first process the leaves, then work our way up to the root.

The key insight is that we don't need to track all node values in a subtree - we only care about the k largest values. Why? Because if a node's value is greater than the k-th largest value in its subtree, it's automatically greater than at least k nodes. This is where the heap data structure becomes useful.

Consider how we can efficiently maintain the k largest values as we traverse the tree. A min-heap of size k is perfect for this - the smallest element in this heap represents the k-th largest value we've seen. When we add a new value:

• If the heap has fewer than k elements, we simply add it
• If the heap already has k elements, we add the new value and remove the smallest, ensuring we always keep the k largest

The clever trick in the solution is using negative values in the heap. Python's heapq is a min-heap, but by negating values when we insert them, the "minimum" negative value actually represents the maximum positive value. So -l[0] gives us the k-th largest value in the subtree.

As we traverse using DFS, at each node we:

1. Merge the heaps from left and right subtrees (keeping only the k largest)
2. Check if we have at least k nodes and if the current node's value beats the k-th largest
3. Add the current node's value to the heap for its parent to use

This bottom-up accumulation ensures that by the time we evaluate any node, we have complete information about its subtree, allowing us to make the "great enough" determination accurately.

Learn more about Tree, Depth-First Search, Divide and Conquer and Binary Tree patterns.

Solution Approach

The solution implements a post-order DFS traversal with a clever heap-based tracking system. Let's break down the implementation:

Core Data Structure - Min-Heap with Size Limit: The push helper function maintains a heap of at most k elements:

def push(pq, x):
    heappush(pq, x)
    if len(pq) > k:
        heappop(pq)

This ensures we only keep the k largest values (stored as negatives) by removing the smallest when the heap exceeds size k.

DFS Traversal Pattern: The dfs function processes nodes in post-order:

def dfs(root):
    if root is None:
        return []
    l, r = dfs(root.left), dfs(root.right)

First recursively process left and right children, getting their heaps of k largest values.

Merging Subtree Information:

for x in r:
    push(l, x)

We merge the right subtree's heap into the left subtree's heap. By reusing l as our accumulator and pushing elements through our size-limited push function, we efficiently maintain only the k largest values from both subtrees.

Checking "Great Enough" Condition:

if len(l) == k and -l[0] < root.val:
    nonlocal ans
    ans += 1

• len(l) == k: Ensures the subtree has at least k nodes
• -l[0] < root.val: Since we store negative values, -l[0] gives us the k-th largest value. If the current node's value is greater than this, it's greater than at least k nodes in its subtree

Adding Current Node:

push(l, -root.val)
return l

We add the current node's negated value to the heap (maintaining the k largest) and return this heap for the parent node to use.

The algorithm's time complexity is O(n * k * log k) where n is the number of nodes, as each node performs heap operations on a heap of size at most k. The space complexity is O(h * k) where h is the height of the tree, due to the recursion stack and heaps maintained at each level.

Example Walkthrough

Let's walk through a concrete example with k = 2 and the following binary tree:

       5
      / \
     3   7
    /   / \
   2   6   9

We'll trace through the DFS traversal to see how the algorithm determines which nodes are "great enough."

Step 1: Process leaf node 2

• DFS returns to node 2
• Left and right children are null, so l = [] and r = []
• After merging: l = [] (empty)
• Check: len(l) = 0 < k, so node 2 is NOT great enough
• Add node 2: push(l, -2) → l = [-2]
• Return [-2] to parent

Step 2: Process node 3

• Left child returns [-2], right child returns []
• After merging: l = [-2]
• Check: len(l) = 1 < k, so node 3 is NOT great enough
• Add node 3: push(l, -3) → l = [-3, -2] (heap maintains smallest at top)
• Return [-3, -2] to parent

Step 3: Process leaf node 6

• Similar to node 2, returns [-6] to parent

Step 4: Process leaf node 9

• Similar to node 2, returns [-9] to parent

Step 5: Process node 7

• Left child returns [-6], right child returns [-9]
• Merge right into left: push([-6], -9) → l = [-9, -6]
• Check: len(l) = 2 = k and -l[0] = -(-9) = 9. Is 7 > 9? No
• Node 7 is NOT great enough (its value isn't greater than at least 2 nodes)
• Add node 7: push(l, -7) → l = [-9, -6, -7]
• Since len(l) > k, pop smallest: l = [-7, -6] (keeping 2 largest: 7 and 6)
• Return [-7, -6] to parent

Step 6: Process root node 5

• Left child returns [-3, -2], right child returns [-7, -6]
• Merge right into left:
  • push([-3, -2], -7) → [-7, -2, -3], pop smallest → [-3, -2]
  • push([-3, -2], -6) → [-6, -2, -3], pop smallest → [-3, -2]
• After merging: l = [-3, -2] (the 2 largest values from all descendants)
• Check: len(l) = 2 = k and -l[0] = -(-3) = 3. Is 5 > 3? Yes!
• Node 5 IS great enough! ans = 1
• Add node 5: push(l, -5) → [-5, -2, -3], pop smallest → [-3, -2]

Final Result: Only node 5 is "great enough" because:

• Its subtree has at least 2 nodes (it has 6 total)
• Its value (5) is greater than at least 2 nodes' values in its subtree (greater than 2 and 3)

The algorithm returns ans = 1.

Solution Implementation

Time and Space Complexity

Time Complexity: O(n * k * log k) where n is the number of nodes in the tree.

The algorithm performs a depth-first search traversal visiting each node exactly once. At each node, we perform the following operations:

• Merge the right subtree's heap into the left subtree's heap: This involves pushing at most k elements from the right heap into the left heap, where each push operation takes O(log k) time.
• Check if the current node satisfies the condition: O(1) time to access the minimum element.
• Push the current node's value into the heap: O(log k) time.

Since we maintain heaps of size at most k, each heap operation (push/pop) takes O(log k) time. In the worst case, at each node we perform O(k) heap operations, resulting in O(k * log k) work per node. With n nodes total, the overall time complexity is O(n * k * log k).

Space Complexity: O(n + h * k) where h is the height of the tree.

The space complexity consists of:

• Recursive call stack: O(h) where h is the height of the tree (could be O(n) in the worst case of a skewed tree).
• Heap storage: At each level of recursion, we maintain a heap of size at most k. In the worst case, we have O(h) heaps active simultaneously in the call stack, each containing at most k elements, giving us O(h * k) space.
• Total nodes stored: Across all heaps during the entire execution, we store at most n node values (each node's value is stored once).

Therefore, the total space complexity is O(n + h * k), which simplifies to O(n) when k is constant and the tree is balanced (h = O(log n)), or O(n * k) in the worst case of a skewed tree where h = O(n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

Pitfall 1: Incorrect Heap Size Check Timing

Problem: A common mistake is checking if a node is "great enough" after adding the current node's value to the heap, rather than before.

Incorrect approach:

# Wrong: Adding node value first, then checking
maintain_k_largest(left_heap, -node.val)
if len(left_heap) == k and -left_heap[0] < node.val:
    count += 1

Why it fails: After adding the current node's value, the heap now includes the node itself when checking if it's greater than k nodes. This leads to comparing the node against itself, which violates the problem's requirement that the node should be greater than k other nodes in its subtree.

Solution: Always check the "great enough" condition before adding the current node's value to the heap:

# Correct: Check first, then add
if len(left_heap) == k and -left_heap[0] < node.val:
    count += 1
maintain_k_largest(left_heap, -node.val)

Pitfall 2: Forgetting to Negate Values for Max-Heap Simulation

Problem: Using positive values in the heap while expecting max-heap behavior from Python's min-heap.

Incorrect approach:

# Wrong: Using positive values in a min-heap
maintain_k_largest(left_heap, node.val)  # This maintains k smallest, not largest
if len(left_heap) == k and left_heap[0] < node.val:  # Wrong comparison
    count += 1

Why it fails: Python's heapq implements a min-heap. Without negating values, you'll track the k smallest values instead of the k largest, completely inverting the logic.

Solution: Consistently negate values when pushing to simulate max-heap behavior:

# Store as negative
maintain_k_largest(left_heap, -node.val)
# Compare with negated heap top
if len(left_heap) == k and -left_heap[0] < node.val:
    count += 1

Pitfall 3: Misunderstanding the Subtree Size Requirement

Problem: Checking len(left_heap) >= k instead of exactly len(left_heap) == k.

Incorrect approach:

# Wrong: Using >= instead of ==
if len(left_heap) >= k and -left_heap[0] < node.val:
    count += 1

Why it fails: Before adding the current node, if the heap has exactly k-1 elements, the subtree (excluding current node) has k-1 nodes. After adding the current node, the total becomes k nodes. Using >= would incorrectly count nodes whose subtrees have fewer than k total nodes.

Solution: Use strict equality to ensure exactly k nodes exist in the subtree when including the current node:

# Correct: Ensures subtree has exactly k nodes after including current
if len(left_heap) == k and -left_heap[0] < node.val:
    count += 1