Count pairs from 2 BST whose sum is equal to given value X

The problem, simply

You get two binary search trees (BSTs), A with m numbers and B with n numbers. Numbers inside one tree are all different. You also get a target x.

Count the pairs (a, b) where a is in A, b is in B, and a + b == x.

A story to remember it

Picture it

Examples

Way 1: the simple way

For every number a in A, search B for x - a. B is a BST, so each search walks one road.

1. Visit every spot of A.
2. Search B for x - a (smaller goes left, bigger goes right).
3. If found, add 1 to the count.

class Solution {
    public int countPairs(TreeNode a, TreeNode b, int x) {
        if (a == null) return 0;
        int here = contains(b, x - a.val) ? 1 : 0;
        return here + countPairs(a.left, b, x) + countPairs(a.right, b, x);
    }

    private boolean contains(TreeNode node, int target) {
        while (node != null) {
            if (target == node.val) return true;
            node = target < node.val ? node.left : node.right;
        }
        return false;
    }
}

How long it takes: m searches, each up to the height of B, O(m * h2). If B is a stick, that is O(m * n). Extra room: the recursion on A, O(h1).

Way 2: the clever way

If both trees were sorted lists, this would be the classic two finger trick: one finger at the small end of A, one at the big end of B. A BST is a sorted list in disguise, so walk it with a pile.

1. Pile A walks A from smallest up: push its left road, and after each pop push the left road of the popped spot's right child.
2. Pile B walks B from biggest down: push its right road, and after each pop push the right road of the popped spot's left child.
3. Take a from A and b from B. If a + b == x, count it and take the next of both. If the sum is too small, take the next a (bigger). If too big, take the next b (smaller).
4. Stop when either pile runs out.

import java.util.*;

class Solution {
    public int countPairs(TreeNode a, TreeNode b, int x) {
        Deque<TreeNode> up = new ArrayDeque<>();    // A, smallest first
        Deque<TreeNode> down = new ArrayDeque<>();  // B, biggest first
        pushLeft(a, up);
        pushRight(b, down);
        int count = 0;
        while (!up.isEmpty() && !down.isEmpty()) {
            int sum = up.peek().val + down.peek().val;
            if (sum == x) {
                count++;
                pushLeft(up.pop().right, up);       // next a
                pushRight(down.pop().left, down);   // next b
            } else if (sum < x) {
                pushLeft(up.pop().right, up);       // need a bigger a
            } else {
                pushRight(down.pop().left, down);   // need a smaller b
            }
        }
        return count;
    }

    private void pushLeft(TreeNode n, Deque<TreeNode> pile) {
        while (n != null) { pile.push(n); n = n.left; }
    }

    private void pushRight(TreeNode n, Deque<TreeNode> pile) {
        while (n != null) { pile.push(n); n = n.right; }
    }
}

How long it takes: every spot of both trees is pushed and popped at most once, O(m + n). Extra room: the two piles, each one road tall, O(h1 + h2).

Watch out for

• Walking both trees the same direction. Then a too-small sum has no move that is sure to help.
• Stopping when one pile is empty is right. With no a left there are no more pairs.
• Sums can overflow int for huge values. Use long sum if the numbers can be near Integer.MAX_VALUE.

What interviewers ask next

• Use a hash set instead. Put all of A in a set, walk B and look up x - b. O(m + n) time, but O(m) extra room.
• Two numbers in one BST that add up to x (LeetCode 653). Same two piles on the same tree, smallest up and biggest down, stopping when they meet.
• Count pairs with sum at most x. When a + b <= x, every number in B up to b also works with a, so this needs sizes or a different sweep. With sorted lists it is the usual two pointer count.
• Three numbers from three trees. Fix each a from the first tree, then run this two pile walk on the other two. O(m * (n + p)).

Remember this

Grown-up words: the seesaw lines are the two pointer technique, walking a BST up or down with a pile is a BST iterator (forward and reverse inorder), and O(h1 + h2) is the height of the two piles.