from collections import defaultdict, deque

from ds import TreeNode

from typing import List

"""

Definition of TreeNode:

class TreeNode:

def __init__(self, val):

self.val = val

self.left = None

self.right = None

"""

def binary_tree_columns(root: TreeNode) -> List[List[int]]:

if not root:

return []

column_map = defaultdict(list)

leftmost_column = rightmost_column = 0

queue = deque([(root, 0)])

while queue:

node, column = queue.popleft()

if node:

# Add the current node's value to its corresponding list in the hash

# map.

column_map[column].append(node.val)

leftmost_column = min(leftmost_column, column)

rightmost_column = max(rightmost_column, column)

# Add the current node's children to the queue with their respective

# column ids.

queue.append((node.left, column - 1))

queue.append((node.right, column + 1))

# Construct the output list by collecting values from each column in the hash

# map in the correct order.

return [column_map[i] for i in range(leftmost_column, rightmost_column + 1)]