Find maximum sum possible equal sum of three stacks

The problem, simply

You get three stacks of positive numbers, given as arrays where index 0 is the top. You may only remove numbers from the top of a stack, as many as you like.

Remove numbers so that all three stacks have the same sum, and that sum is as big as possible. Give back that sum. If the only way is to empty them, the answer is 0.

A story to remember it

Picture it

Examples

Way 1: the simple way

Write down every sum each stack can still have, then look for the biggest sum that all three share.

1. For each stack, compute the sum of everything from index i down to the bottom, for every i. These suffix sums are all the sums that stack can reach.
2. Try every triple of choices and keep the biggest sum where all three are equal.

class Solution {
    public long maxEqualSum(int[] s1, int[] s2, int[] s3) {
        long[] a = suffix(s1), b = suffix(s2), c = suffix(s3);
        long best = 0;
        for (long x : a)
            for (long y : b)
                for (long z : c)
                    if (x == y && y == z) best = Math.max(best, x);
        return best;
    }

    // suf[i] = sum of s[i..end], the sum left after removing the top i numbers
    private long[] suffix(int[] s) {
        long[] suf = new long[s.length + 1];
        for (int i = s.length - 1; i >= 0; i--) suf[i] = suf[i + 1] + s[i];
        return suf;
    }
}

How long it takes: three nested loops, O(n1 * n2 * n3). Extra room: the three suffix arrays, O(n1 + n2 + n3).

Way 2: the clever way

Do what the kids did.

1. Add up each stack.
2. While the three sums are not equal, find the stack with the biggest sum and remove its top number. Move that stack's pointer down by one.
3. If any stack runs out, the answer is 0. Otherwise the shared sum is the answer.

Why lift from the biggest? The final shared sum can never be more than the smallest current sum. The biggest stack must come down anyway, and lifting its top block is the smallest step it can take. It never jumps past a sum that the others could still reach.

class Solution {
    public long maxEqualSum(int[] s1, int[] s2, int[] s3) {
        long a = 0, b = 0, c = 0;
        for (int x : s1) a += x;
        for (int x : s2) b += x;
        for (int x : s3) c += x;

        int i = 0, j = 0, k = 0;                 // how many tops each stack has lost
        while (!(a == b && b == c)) {
            if (a >= b && a >= c) {
                if (i == s1.length) return 0;
                a -= s1[i++];                    // lift from the heaviest stack
            } else if (b >= a && b >= c) {
                if (j == s2.length) return 0;
                b -= s2[j++];
            } else {
                if (k == s3.length) return 0;
                c -= s3[k++];
            }
        }
        return a;
    }
}

How long it takes: every number is lifted at most once, O(n1 + n2 + n3). Extra room: a few counters, O(1).

Watch out for

• Sums can pass the int range when there are many big numbers. Use long.
• The loop must stop when all three are equal, including when all are 0.
• Ties for the biggest sum: lifting from any of the tied stacks is fine.

What interviewers ask next

• Why does lifting from the biggest never skip the answer? The answer is at most the smallest sum right now. The biggest stack is above that, so it must lose its top block in any solution. Numbers are positive, so lifting it can only bring it closer.
• What if numbers can be negative? Then sums can go up when you lift, and the greedy breaks. Put each stack's suffix sums in a hash set and look for the largest value in all three sets.
• What about k stacks? Same idea with a max heap of sums, or note that the answer is the largest value that appears in every stack's suffix sum set.
• What if you may remove from the bottom instead? Use prefix sums from the other end; the greedy is the same.

Remember this

Grown-up words: the towers are stacks, lifting the top block is a pop, the three counters are three pointers, and the sums a stack can still reach are its suffix sums.