Problem statement

https://leetcode.com/problems/n-ary-tree-level-order-traversal/

Solution

Nothing very special about this problem, in fact it is very similar to problem 0102. Binary Tree Level Order Traversal. The idea is to use any graph traversal, here I use bfs. We keep queue Q with value in current level. Then we do len(Q) extractions from the left and add all children to the right. Also we add all values fro this level to ans. Note that we always keep invariant: first we traverse full first layer, then it is fully removed and second layer added and so on.

Complexity

Time complexity is O(n) to traverse our tree, space complexity is O(n) as well to return answer.

Code

class Solution:
    def levelOrder(self, root):
        if not root: return []
        Q, ans = deque([root]), []
        
        while Q:
            level = []
            for i in range(len(Q)):
                node = Q.popleft()
                for child in node.children:
                    Q.append(child)
                level += [node.val] 
            ans += [level]

        return ans