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])
My answer works but I hate how I implemented it. I am having trouble keeping track of the height of each node


You can use a queue and implement a BFS traversal that follows these conditions:
 The queue will only contain nodes in the same level.
 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