Is there a O(n) solution which does not bottom-up depth on this question? I got asked this question in a google on-stie interview, but I don't think it is possible to solve this question by not using the bottom-up depth and dfs.

Recursively delete leave nodes in a multi-tree. Require O(n) solution

Example input:

      1
   /   |  \
  2    5   3
 / \   | 
7   4  9
    |
    8

should return: [7, 8, 9, 3, 4, 5, 2, 1]

My solution is to use bottom-up depth information to hold nodes with the same depth and then concatenate them. It is O(n) in time time complexity and O(n) in space complexity.

def remove_leaves(self, root):
    nodes_with_the_same_depth = []

    def dfs(node):
        if not node:
            return 0

        current_height = 1
        for child in node.children:
            current_height = max(current_height, dfs(child) + 1)

        if len(nodes_with_the_same_depth) == current_height - 1:
            nodes_with_the_same_depth.append([])
        nodes_with_the_same_depth[current_height-1].append(node.val)

        return current_height

    dfs(root)
    ans = []
    for nodes in nodes_with_the_same_depth:
        ans.extend(nodes)

    return ans
	

With the example given above, nodes_with_the_same_depth will be [[7, 8, 9, 3], [4, 5], [2], [1]]. By concatenating sub-arrays, we will get the final answer.

My argument why this can't be solved without using bottom-up depth is that whether it is dfs or bfs, the order of visiting each node is not guarranted to be correct as the requirement. But if we have arrays to store the nodes with the same depth, we are able to put the node in the right array with the same bottom-up tree depth. The bottom-up depth of a node determines the order when a node should be deleted. The interviewer insisted that this could be improved by not introducing the depth, but instead somehow by disconnecting leave nodes from their parents when they are deleted. However, I still don't agree with her since if not knowing the depth, there is not way to make the order of traversal correct.

The above solution using bottom-up tree depth arrays is a valid one but is not what exactly in the interviewer's mind. And the interviewer did not give me enough hint to get closer to what's in her mind which is quite disappointing.

Can anyone prove if I am wrong?