Print all K Sum paths in a Binary tree

The problem, simply

Every node of a binary tree holds a number, and numbers can be negative. A path here means a chain of nodes going straight down: it can start at any node and stop at any node below it, but it never goes up or sideways.

Print every path whose numbers add up to k.

A story to remember it

Picture it

Examples

Way 1: the simple way

Try every starting node. From each one, walk down every way, adding numbers as you go, and print every time the sum equals k. Do not stop when you hit k, because a minus number further down can bring you back to k again.

1. For each node s in the tree, start a walk at s with sum 0.
2. In the walk, add the node's value. If the sum is k, save the path so far.
3. Keep walking into both children.

import java.util.*;

class Solution {
    private List<List<Integer>> found = new ArrayList<>();

    public List<List<Integer>> kSumPaths(TreeNode root, int k) {
        found.clear();
        startEverywhere(root, k);
        return found;
    }

    private void startEverywhere(TreeNode node, int k) {
        if (node == null) return;
        walkDown(node, 0, k, new ArrayList<>());   // all paths that start here
        startEverywhere(node.left, k);
        startEverywhere(node.right, k);
    }

    private void walkDown(TreeNode node, int sum, int k, List<Integer> path) {
        if (node == null) return;
        sum += node.val;
        path.add(node.val);
        if (sum == k) found.add(new ArrayList<>(path));   // do not stop, keep going
        walkDown(node.left, sum, k, path);
        walkDown(node.right, sum, k, path);
        path.remove(path.size() - 1);
    }
}

How long it takes: each node is walked over once for every ancestor that starts a walk, so O(n * h). That is O(n log n) for a bushy tree and O(n^2) for a stick. Copying each answer costs extra on top. Extra room: the path and the call stack, O(h).

Way 2: the clever way

Use the notebook from the hike. Keep the running total from the root, and a map from "total" to the spots on the current path where that total happened. Spot 0 means "before the root", with total 0.

1. Step onto a node. The new total is total = old total + value.
2. Ask the map for total - k. For every spot j it lists, the nodes after spot j up to here add up to exactly k. Save them.
3. Write total into the map at the current spot.
4. Go into the left child and the right child.
5. On the way back up, erase total from the map, so other forks do not see it.

import java.util.*;

class Solution {
    private List<List<Integer>> found = new ArrayList<>();
    private List<Integer> path = new ArrayList<>();               // values from the root down
    private Map<Integer, List<Integer>> seen = new HashMap<>();   // total -> spots where it happened

    public List<List<Integer>> kSumPaths(TreeNode root, int k) {
        found.clear();
        path.clear();
        seen.clear();
        seen.computeIfAbsent(0, x -> new ArrayList<>()).add(0);   // before the root, total 0
        go(root, 0, k);
        return found;
    }

    private void go(TreeNode node, int total, int k) {
        if (node == null) return;
        path.add(node.val);
        total += node.val;
        for (int spot : seen.getOrDefault(total - k, Collections.emptyList())) {
            found.add(new ArrayList<>(path.subList(spot, path.size())));   // the nodes after that spot
        }
        seen.computeIfAbsent(total, x -> new ArrayList<>()).add(path.size());
        go(node.left, total, k);
        go(node.right, total, k);
        List<Integer> spots = seen.get(total);
        spots.remove(spots.size() - 1);    // walking back up: erase this total
        path.remove(path.size() - 1);
    }
}

How long it takes: one visit per node and one map question per node, so O(n) to find all the paths, plus the time to copy out each answer. Extra room: the map and the path only hold the current way down, O(h).

Watch out for

• Ask the map before you write the current total into it. Otherwise, with k = 0, every node would wrongly match an empty path ending at itself.
• Forgetting to erase on the way back up lets totals from the left fork leak into the right fork and invents paths that are not straight.
• Start the map with total 0 at spot 0, or paths that start at the root are missed.

What interviewers ask next

• Only count the paths (LeetCode "path sum III"). Keep a map from total to how many times it happened instead of the spots: O(n) time.
• The textbook way. Keep the current path and, at each node, add values backwards from the node to the root, printing whenever the sum is k. It is one walk, but still O(n * h), like Way 1.
• Paths may go up and then down (any two nodes). That is "maximum path sum" style: each node combines the best from its left and right. For counting all such paths, it gets much harder.
• Root to leaf paths only. Then a simple walk with a running sum works, and you only check at leaves.

Remember this

Grown-up words: the running totals are prefix sums, the notebook is a hash map, erasing on the way back is backtracking, and the same trick on an array is "subarray sum equals k".