Merge neighbor sums from many binary trees into the last: Frontier Interview's solution

Question: Merge neighbor sums from many binary trees into the last

Companies: Snowflake

Difficulty: Medium

Problem statement

You are given a list trees of binary trees (each given by its root, possibly empty). The neighbor sum of a node is its own value plus the values of its left and right children, where a missing child counts as 0.

Treat every tree except the last one as a source. The last tree is the target. Process the sources from left to right. For each node in a source, find the node at the same position in the target. A position is the path from the root: the root, then left or right at each step.

If the target has a node there, add the source node's neighbor sum to its value.
If the target has no node there, create one whose value is the neighbor sum, and attach it at that position.

Always compute neighbor sums from the source trees' original values. Do not modify the source trees. Return the root of the updated target. Return None if trees is empty, or if every tree is empty.

def copy_neighbor_sums(trees: list[Optional[TreeNode]]) -> Optional[TreeNode]: ...

Constraints

0 <= len(trees) <= 10^3
Up to 10^5 nodes in total; node values are integers (they may be negative)
Trees can be very unbalanced (depth up to 10^5)

Example 1 (trees written in level order, None = missing)

trees = [[2, 4, 1], [10, 1]]

The source root's neighbor sum is 2 + 4 + 1 = 7, so the target root becomes 17. The left node's neighbor sum is 4, so the target's left node becomes 5. The target has no right child, so a node with value 1 is created.
Result: [17, 5, 1].

Example 2

trees = [[1, 2], [3], [0]]

Each source adds 3 to the root (1 + 2 + 0, then 3 + 0 + 0), so the root becomes 6. The first source's left child creates a node with value 2.
Result: [6, 2].

Reconstructed from the reported question; details may differ from what interviewers ask.

Approach

This is LeetCode 617 (Merge Two Binary Trees) with two changes. There are many
trees instead of two, and a source node does not add its own value. It adds its
neighbor sum: its own value plus the values of its left and right children,
where a missing child counts as 0.

"Position" means the path from the root (root, root.left, root.right.left, and
so on). Two nodes in different trees correspond when they sit at the same path.

Brute force. Turn every tree into a map from path to node, for example
"LRL" -> node. Then compute each source node's neighbor sum, and for each path
either add it to the target node or build the missing chain of nodes. This
works, but building paths costs O(depth) per node, and you have to rebuild
parents in the right order. It also hides the simple recursive structure.

Optimal: walk the source and the target together. For each source tree in
order, run one traversal that holds a pair (src, dst) at the same position.

If src is None, nothing at or below this position exists in the source.
Return dst unchanged.
Compute ns = src.val + val(src.left) + val(src.right). Read these values
from the source tree, which we never change.
If dst is None, create a node with value ns. Otherwise add ns to
dst.val.
Recurse into the children: dst.left = walk(src.left, dst.left), and the
same for the right side. The return value matters. It is how a newly created
child is attached to its parent.

Because the walk goes top-down, a parent is always present (or created) before
its children. So a deep source node never needs special handling for "missing
parents" in the target.

Invariants.

Source trees are read-only. Every neighbor sum uses the source's original
values, never a value that an earlier merge changed.
After processing source k, each target node holds its original value
(0 if it did not exist) plus the neighbor sums of sources 0..k at its
position.
Addition commutes, so the final values do not depend on the order of the
sources. The order is in the statement only to define the process.

Use an explicit stack instead of recursion, because a skewed tree with 10^5
nodes would overflow Python's default recursion limit. The stack holds the
source node, the parent in the target, and the side to attach on.

Code

from __future__ import annotations
from dataclasses import dataclass
from typing import Optional


@dataclass
class TreeNode:
    val: int
    left: Optional["TreeNode"] = None
    right: Optional["TreeNode"] = None


def _val(node: Optional[TreeNode]) -> int:
    return node.val if node else 0


def _neighbor_sum(node: TreeNode) -> int:
    return node.val + _val(node.left) + _val(node.right)


def _apply(src: Optional[TreeNode], dst: Optional[TreeNode]) -> Optional[TreeNode]:
    """Add src's neighbor sums onto dst in place; return dst's (possibly new) root."""
    if src is None:
        return dst
    if dst is None:
        dst = TreeNode(_neighbor_sum(src))
    else:
        dst.val += _neighbor_sum(src)
    # Each entry: (source child, target parent, side). Children are attached as we go,
    # so a parent always exists before its children are visited.
    stack: list[tuple[TreeNode, TreeNode, str]] = []
    for side in ("left", "right"):
        child = getattr(src, side)
        if child is not None:
            stack.append((child, dst, side))
    while stack:
        s, parent, side = stack.pop()
        d = getattr(parent, side)
        if d is None:
            d = TreeNode(_neighbor_sum(s))
            setattr(parent, side, d)
        else:
            d.val += _neighbor_sum(s)
        for child_side in ("left", "right"):
            child = getattr(s, child_side)
            if child is not None:
                stack.append((child, d, child_side))
    return dst


def copy_neighbor_sums(trees: list[Optional[TreeNode]]) -> Optional[TreeNode]:
    """Apply every tree's neighbor sums (all but the last) onto the last tree."""
    if not trees:
        return None
    target = trees[-1]
    for src in trees[:-1]:
        target = _apply(src, target)
    return target


# ---- helpers for tests: level-order lists with None holes ----
def build(values: list[Optional[int]]) -> Optional[TreeNode]:
    if not values or values[0] is None:
        return None
    root = TreeNode(values[0])
    queue, i = [root], 1
    for node in queue:
        if i >= len(values):
            break
        for side in ("left", "right"):
            if i < len(values) and values[i] is not None:
                child = TreeNode(values[i])
                setattr(node, side, child)
                queue.append(child)
            i += 1
    return root


def dump(root: Optional[TreeNode]) -> list[Optional[int]]:
    out: list[Optional[int]] = []
    queue: list[Optional[TreeNode]] = [root]
    for node in queue:
        if node is None:
            out.append(None)
            continue
        out.append(node.val)
        queue.extend([node.left, node.right])
    while out and out[-1] is None:
        out.pop()
    return out


if __name__ == "__main__":
    a = build([2, 4, 1])
    assert dump(copy_neighbor_sums([a, build([10, 1])])) == [17, 5, 1]
    assert dump(a) == [2, 4, 1]  # source untouched
    assert dump(copy_neighbor_sums([build([1, None, 5, 3]), None])) == [6, None, 8, 3]
    assert dump(copy_neighbor_sums([build([1, 2]), build([3]), build([0])])) == [6, 2]
    assert copy_neighbor_sums([]) is None
    assert copy_neighbor_sums([None, None]) is None
    assert dump(copy_neighbor_sums([build([7])])) == [7]
    print("ok")

Complexity

Let N be the total number of nodes across all source trees and M the number
of nodes in the last tree.

Time: O(N). Each source node is visited once, and its neighbor sum takes
O(1). Target nodes that no source reaches are never touched.
Extra space: O(H) for the stack, where H is the height of the current
source tree (O(n) for a skewed tree, O(log n) for a balanced one). The result
is built in place in the last tree. New nodes appear only at positions that
exist in some source and not in the target, so the output has at most M + N
nodes.

Test cases

Trees are written in level order, with None for a missing node.

[[2, 4, 1], [10, 1]] returns [17, 5, 1]. The root gets 2+4+1 = 7, the left
node gets 4, and the right node is created with 1.
[[1, None, 5, 3], None] returns [6, None, 8, 3]. The target is empty, so
every node is created with its neighbor sum (root 1+0+5, right 5+3+0, leaf 3).
[[1, 2], [3], [0]] returns [6, 2]. The root gets 3 from the first source
and 3 from the second, and the left node is created with 2.
[] returns None.
[None, None] returns None. There is nothing to apply and nothing to
return.
[[7]] returns [7]. With a single tree, there are no sources.
Negative values: [[-1, 2, -3], [0]] returns [-2, 2, -3].
A source tree is never mutated. After the call, [2, 4, 1] still reads
[2, 4, 1].
A deep skewed source (10^5 nodes down the left side) finishes without a
recursion error because the traversal is iterative.

Follow-ups

Return a new tree instead of mutating the last one. Clone the target
lazily: create a fresh node at every position the walk visits, and share any
untouched target subtree by reference (or deep-copy it if sharing is not
allowed). The time is the same. The extra space becomes O(N + M) if you copy
everything.
Process the sources in parallel or out of order. Addition commutes, so
the order does not matter. Fold every source into one "delta" tree first,
then apply that once. Or merge sources in pairs like a tournament, then
apply the result to the target.
Neighbors include the parent too (or the definition changes, for
example to the subtree sum). Pass the parent's value down in the stack
entry. For subtree sums, run a post-order pass on each source first to
compute the sums, then do the same top-down merge.
N-ary trees, or trees given as arrays or serialized strings. Use the
same paired walk, matching children by index. For level-order arrays, a
position is an index (2i+1, 2i+2). Pad the target array and add the
values with no pointers at all, but watch the memory blowup on sparse,
deep trees.
Very large trees that do not fit in memory, or streaming sources. Key
each node by its path, for example as a bit string, and emit (path, neighbor_sum) records. Then group by key and sum, MapReduce style. This
ties into Snowflake's focus on data engineering.

Interviewer expectations

Mid-level (L4)

Recognizes the link to Merge Two Binary Trees and defines "position" as the
path from the root.
Writes a correct recursive solution that attaches newly created children
through return values.
Handles an empty list, an empty target, and a single tree.

Senior (L5)

Reads neighbor sums from unmodified source values and says so. Notes that
the order does not matter because addition commutes.
Points out the recursion-depth risk on skewed trees and switches to an
explicit stack without being asked.
Gives exact complexity in terms of the source sizes, not the target size, and
writes useful tests including negatives and mutation checks.

Staff (L6+)

Discusses mutation versus a persistent or copy-on-write result, and the
trade-offs of sharing structure.
Generalizes the design: one pass that builds a delta tree, a pluggable
"contribution" function, and N-ary or array forms.
Extends to scale: path-keyed records, aggregation in the style of a group-by,
and parallel reduction across many sources.