My answer works but I hate how I implemented it. I am having trouble keeping track of the height of each node


  • 0
    G
    class Solution(object):
        def levelOrder(self, root):
            """
            :type root: TreeNode
            :rtype: List[List[int]]
            """
            if not root:
                return []
            ret = [[root.val, 1]]
            stack = [[root, 1]]
            while stack:
                node, depth = stack.pop()
                if not node:
                    continue
                if node.left and node.right:
                    ret.append([node.left.val, node.right.val, depth+1])
                elif node.left and not node.right:
                    ret.append([node.left.val, depth+1])
                elif node.right and not node.left:
                    ret.append([node.right.val, depth+1])
                stack.append([node.right, depth+1])
                stack.append([node.left, depth+1])
            
            ret.sort(self.sort_by_depth)  # sort all elements by their depth
            final = []
            i = 0
            temp = ret[0][:-1]
            while i <= len(ret)-2:  # merge the elements that are on the same level
                if ret[i][-1] != ret[i+1][-1]:
                    final.append(temp)
                    temp = ret[i+1][:-1]
                else:
                    temp = temp+ret[i+1][:-1]
                i += 1
            final.append(temp)
            return final
            
        def sort_by_depth(self, a, b):
            return cmp(a[-1], b[-1])

  • 0
    N

    You can use a queue and implement a BFS traversal that follows these conditions:

    1. The queue will only contain nodes in the same level.
    2. The nodes will be inserted to the queue from left to right.

    That way you won't have to keep track of the level of the node before adding it to the result list:

    from collections import deque
    class Solution(object):
        def levelOrder(self, root):
            if not root:
                return []
    
            res, queue = [], deque([root])
            while queue:
    
                # Deque all children in the same level.
                children = []
                while queue:
                    children.append(queue.popleft())
    
                # Add the children to result
                res.append([x.val for x in children])
    
                # Insert children of children in left to right order
                for child in children:
                    for x in [y for y in [child.left, child.right] if y]:
                        queue.append(x)
    
            return res
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.