Leetcode Daily Question 23/10/2024 - 2641. Cousins in Binary Tree II

Problem Description:

Given the root of a binary tree, replace the value of each node in the tree with the sum of all its cousins' values. | ------------------------------------------------------------ | Two nodes of a binary tree are cousins if they have the same depth with different parents. | ------------------------------------------------------------ | Return the root of the modified tree. | ------------------------------------------------------------ | Note that the depth of a node is the number of edges in the path from the root node to it.

from typing import Optional
class TreeNode:
  def __init__(self, val = 0, left = None, right = None):
    self.val = val
    self.left = left
    self.right = right
    

class Solution:
  def replaceValueInTree(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
    if not root:
      return None
    
    queue = deque([root])
    level_sums = []
    
    while queue:
      level_size = len(queue)
      level_sum = 0
      for _ in range(level_size):
        node = queue.popleft()
        level_sum += node.val
        if node.left:
          queue.append(node.left)
        if node.right:
          queue.append(node.right)
      level_sums.append(level_sum)
      
      
    depth = 1
    nqueue = deque([root])
    root.val = 0
    
    while nqueue:
      level_size = len(nqueue)
      for _ in range(nqueue):
        node = nqueue.popleft()
        sib_sum = (node.left.val if node.left else 0) + (node.right.val if node.right else 0)
        if node.left:
          node.left.val = level_sums[depth] - sib_sum
          nqueue.append(node.left)
        if node.right:
          node.right.val = level_sums[depth] - sib_sum
          nqueue.append(node.right)
      depth += 1
        
    return root

Explanation:

The idea is use BFS twice. Once to calculate the level sum for each level of the tree. Once to update the values of the nodes base on its siblings (nodes with same parents).

Lets start by calculating the level sums.

queue = deque([root])
level_sums = []

while queue:
  level_size = len(queue)
  level_sum = 0
  for _ in range(level_size):
    node = queue.popleft()
    level_sum += node.val
    if node.left:
      queue.append(node.left)
    if node.right:
      queue.append(node.right)
  level_sums.append(level_sum)

Since every node can only have a maximum of two children, it follows that every node can have at most one sibling.

So if we know the level sum, we can calculate the value of the node by subtracting the sibling sum (the node itself + its sibling if exists) from the level sum.

nqueue = deque([root])
depth = 1 # We are adjusting the value of children, so we need the next level.
root.val = 0 # Root has no parent, so it cannot have any siblings / cousins.

while nqueue:
  level_size = len(nqueue)
  for _ in range(level_size):
    node = nqueue.popleft()
    sib_sum = (node.left.val if node.left else 0) + (node.right.val if node.right else 0)
    if node.left:
      node.left.val = level_sums[depth] - sib_sum
      nqueue.append(node.left)
    if node.right:
      node.right.val = level_sums[depth] - sib_sum
      nqueue.append(node.right)
  depth += 1

return root

The time complexity is O(n). We visit each node twice with the two bfs traversal. So time complexity is O(2n) ≈ O(n)

The space complexity is O(n), as we store the sum of each level. There can be at most n levels for n nodes.

The full solution can be found near the top of the page.