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

• ``````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])``````

• 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
``````

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