2641.cousins-in-binary-tree-ii.md

2641. Cousins in Binary Tree II

BFS

We traverse the tree level by level, and sum up the values of the left and right child at the next level. Then we update the child nodes with total sum - (left + right) node values.

   A       B     C    // Current traversal level
 /   \   /   \     \
1     5 2     0    10 = 18

// Update the node.
 /   \   /   \     \
12   12 16   16     8

// Why 12? 18 - (1 + 5) = 12
// Why 16? 18 - (2 + 0) = 16
// Why 8? 18 - 10 = 8

fun replaceValueInTree(root: TreeNode?): TreeNode? {
    if (root == null) return null
    // Special case, we update the root directly.
    root.`val` = 0
    val queue = ArrayDeque<TreeNode>()
    queue.addLast(root)
    while (queue.isNotEmpty()) {
        val size = queue.size
        var sum = 0
        repeat(size) {
            val node = queue.removeFirst()
            sum += (node.left?.`val` ?: 0) + (node.right?.`val` ?: 0)
            queue.addLast(node)
        }
        repeat(size) {
            val node = queue.removeFirst()
            var leftRightSum = (node.left?.`val` ?: 0) + (node.right?.`val` ?: 0)
            if (node.left != null) {
                node.left.`val` = sum - leftRightSum
                queue.addLast(node.left)
            }
            if (node.right != null) {
                node.right.`val` = sum - leftRightSum
                queue.addLast(node.right)
            }
        }
    }
    return root
}

• Time Complexity: O(n).
• Space Complexity: O(n).