Bishu and Soldiers

The problem, simply

There are n soldiers, each with a power number. Bishu fights them. With power X he beats every soldier whose power is at most X.

You get q questions. Each question gives a value of X. For each one, answer two numbers: how many soldiers Bishu beats, and the sum of their powers. There can be a hundred thousand soldiers and a hundred thousand questions, so each answer must be fast.

A story to remember it

Picture it

Examples

Way 1: the simple way

For each question, walk over every soldier. Count the ones with power at most X and add up their powers.

import java.util.*;

class Solution {
    // returns {count, sum} for every question
    public long[][] answer(int[] powers, int[] questions) {
        long[][] out = new long[questions.length][2];
        for (int k = 0; k < questions.length; k++) {
            long count = 0, sum = 0;
            for (int p : powers) {
                if (p <= questions[k]) {   // Bishu wins this fight
                    count++;
                    sum += p;
                }
            }
            out[k][0] = count;
            out[k][1] = sum;
        }
        return out;
    }
}

How long it takes: every question walks every soldier, O(n * q). With a hundred thousand of each that is ten billion steps. Extra room: O(1) besides the answers.

Way 2: the clever way

Do the slow work once, then answer each question with a few jumps.

1. Sort the powers from weakest to strongest.
2. Build running totals: pre[i] is the sum of the first i sorted powers, and pre[0] = 0.
3. For a question X, use binary search to find the first soldier stronger than X. Call that spot k. Everyone before it is beaten.
4. The answer is k soldiers with total power pre[k].

import java.util.*;

class Solution {
    // index of the first value bigger than x (the count of values <= x)
    int firstBigger(int[] a, int x) {
        int lo = 0, hi = a.length;           // hi is just past the end
        while (lo < hi) {
            int mid = lo + (hi - lo) / 2;
            if (a[mid] <= x) lo = mid + 1;   // mid is beaten, look right
            else hi = mid;                   // mid is too strong, maybe the first one
        }
        return lo;
    }

    public long[][] answer(int[] powers, int[] questions) {
        int[] a = powers.clone();
        Arrays.sort(a);
        long[] pre = new long[a.length + 1];
        for (int i = 0; i < a.length; i++) pre[i + 1] = pre[i] + a[i];

        long[][] out = new long[questions.length][2];
        for (int k = 0; k < questions.length; k++) {
            int beaten = firstBigger(a, questions[k]);
            out[k][0] = beaten;
            out[k][1] = pre[beaten];         // total power of the first beaten soldiers
        }
        return out;
    }
}

How long it takes: sorting costs O(n log n) once. Each question costs one binary search, O(log n). Together O((n + q) log n). Extra room: the running totals, O(n).

Watch out for

• The sum can get big. A hundred thousand soldiers with power up to a million add up to 10^11, far past what an int holds. Keep sums in long.
• Use "at most X" (<=). Equal power still wins.
• Let hi start at n, just past the end, so "everyone is beaten" can be the answer.
• In Java, read input with a fast reader. A slow Scanner can time out with this many numbers.

What interviewers ask next

• Can a library do the search? Arrays.binarySearch returns any matching spot, not the first or last, so write the upper bound yourself. Another way: a TreeMap from power to (count, sum) so far, and floorEntry(X) gives both answers.
• What if powers change between questions? Use a Fenwick tree (a binary indexed tree) over power values to keep counts and sums that can be updated in O(log n).
• What if each question asks for a range, powers between L and R? Two binary searches, then pre[right] - pre[left].
• What if the questions come all at once? Sort the questions too and walk both lists with two pointers, O(n log n + q log q).

Remember this

Grown-up words: the line from shortest to tallest is a sorted array, the stickers are prefix sums, finding the first one too strong is an upper bound search, and doing work once to answer many questions fast is preprocessing.