Got this problem in Snowflake interview:

Given a graph of height >= k rooted at node = 0, return the minimum set of nodes to delete (not count) such that the resulting tree has height of at most k. To delete a node means to remove it from consideration in the height count, but everything in its subtree is still included.

Here was my initial code using a top-down greedy approach, which works for deleting smallest amount of nodes:

def minimumDeletedNodesForHeightK(n, edges, k):
  graph = [[] for _ in range(n)]
  for a, b in edges:
    graph[a].append(b)
    graph[b].append(a)

  heights = [0] * n
  def precomputeHeights(parent, node):
    heights[node] = 1
    for child in graph[node]:
      if child != parent:
        heights[node] = max(heights[node], precomputeHeights(node, child) + 1)

    return heights[node]
  
  precomputeHeights(-1, 0)
  deleted_nodes = []
  def deleteNodes(parent, node):
    if heights[node] > k:
      deleted_nodes.append(node)
      for child in graph[node]:
        if child != parent:
          deleteNodes(node, child)

  deleteNodes(-1, 0)
  return deleted_nodes

But then I had a follow-up constraint to not just minimize the amount of deleted nodes, but also maximize the score of the deleted nodes (essentially minimizing the sum of their heights and prioritizing nodes lower in the tree).

I am wondering if there is a greedy O(n) solution where if there are multiple children who cause the current node to be > k we must delete the current node, otherwise there exists a better node we can delete from the single child if there is only one child causing that.

Here are test cases for this variation:

Case 1:
Input: n = 5 edges = [[0, 1], [1, 2], [2, 3], [3, 4]] k = 3
Expected: Min Size: 2 Min Score: 3 (Nodes 3 and 4)

Case 2:
Input: n = 7 edges = [[0, 1], [0, 2], [1, 3], [1, 4], [2, 5], [2, 6]] k = 2
Expected: Min Size: 1 Min Score: 3 (Node 0)

Case 3:
Input: n = 7 edges = [[0, 1], [0, 2], [0, 3], [3, 4], [4, 5], [5, 6]] k = 2
Expected: Min Size: 3 Min Score: 6 (Nodes 4, 5, 6)

Case 4:
Input: n = 6 edges = [[0, 1], [1, 2], [2, 3], [0, 4], [4, 5]] k = 3
Expected: Min Size: 1 Min Score: 2 (Node 3)

Would appreciate any help with this, thanks!