Kth smallest number again

The problem, simply

You get n ranges [A, B]. Each range puts every whole number from A to B into one big bag. A number that appears in several ranges goes in only once.

Then you get q questions. Each question gives K: if the bag were sorted, which number would be K-th (counting from 1)? If the bag has fewer than K numbers, answer -1. Ranges can reach 10^18, so you cannot list the numbers. Below, S is how many numbers the ranges hold in total.

A story to remember it

Picture it

Examples

Way 1: the simple way

Put every number from every range into a sorted set, which drops repeats for you. Copy it into a list. Then each question is one lookup.

import java.util.*;

class Solution {
    public long[] kthSmallest(long[][] ranges, long[] queries) {
        TreeSet<Long> bag = new TreeSet<>();
        for (long[] r : ranges)
            for (long x = r[0]; x <= r[1]; x++) bag.add(x);   // every number, once
        List<Long> sorted = new ArrayList<>(bag);
        long[] out = new long[queries.length];
        for (int i = 0; i < queries.length; i++) {
            long k = queries[i];
            out[i] = k > sorted.size() ? -1 : sorted.get((int) (k - 1));
        }
        return out;
    }
}

How long it takes: putting S numbers into a sorted set costs O(S log S), then each question is O(1). Extra room: O(S). With ranges as wide as 10^18, this never finishes.

Way 2: the clever way

Work with whole pieces instead of single numbers.

1. Sort the ranges by start.
2. Merge them: walk left to right, and glue a range onto the last piece if it starts inside it (or right after it). Otherwise start a new piece. Now the pieces do not overlap.
3. Running totals: total[i] is how many numbers the pieces 0 to i hold together.
4. For a question K: if K is more than the last total, answer -1. Otherwise binary search for the first piece whose total reaches K. Say before numbers come before it; the answer is start + (K - before) - 1.

import java.util.*;

class Solution {
    public long[] kthSmallest(long[][] ranges, long[] queries) {
        long[][] r = ranges.clone();
        Arrays.sort(r, (x, y) -> Long.compare(x[0], y[0]));     // by start

        List<long[]> pieces = new ArrayList<>();
        for (long[] cur : r) {
            long[] last = pieces.isEmpty() ? null : pieces.get(pieces.size() - 1);
            if (last == null || cur[0] > last[1] + 1) {
                pieces.add(new long[] { cur[0], cur[1] });       // a gap: new piece
            } else {
                last[1] = Math.max(last[1], cur[1]);             // overlap: glue on
            }
        }

        int p = pieces.size();
        long[] total = new long[p];
        for (int i = 0; i < p; i++) {
            long size = pieces.get(i)[1] - pieces.get(i)[0] + 1;
            total[i] = (i == 0 ? 0 : total[i - 1]) + size;
        }

        long[] out = new long[queries.length];
        for (int q = 0; q < queries.length; q++) {
            long k = queries[q];
            if (p == 0 || k > total[p - 1]) { out[q] = -1; continue; }
            int lo = 0, hi = p - 1;
            while (lo < hi) {                                   // first total that reaches k
                int mid = lo + (hi - lo) / 2;
                if (total[mid] >= k) hi = mid;
                else lo = mid + 1;
            }
            long before = lo == 0 ? 0 : total[lo - 1];
            out[q] = pieces.get(lo)[0] + (k - before) - 1;
        }
        return out;
    }
}

How long it takes: sorting and merging cost O(n log n). Each question is one binary search over the pieces, O(log n). Total O(n log n + q log n). Extra room: the pieces and totals, O(n).

Watch out for

• Everything is huge: starts, ends and K go up to 10^18. Use long everywhere, and do not build arrays the size of a range.
• Off by one inside a piece: the first number of a piece is number 1 of that piece, so the answer is start + (K - before) - 1.
• Gluing touching ranges like [1, 3] and [4, 6] is optional (the count is the same either way), but gluing ranges with a real gap, like [1, 3] and [5, 6], is a bug.
• K larger than the total must print -1, not crash.

What interviewers ask next

• Can you answer each question in O(1)? Not directly, but if the questions come all at once, sort them and walk the pieces with one finger, O(n log n + q log q).
• What if ranges are added between questions? Keep the pieces in a TreeMap from start to end and merge each new range as it arrives. Rebuilding the totals costs O(n) per insert; a balanced tree that stores piece sizes in its nodes brings inserts and questions down to O(log n).
• Count of distinct numbers covered by all ranges? It is the last running total.
• This is "Merge Intervals" (LeetCode 56) followed by a search on prefix sums. Both pieces show up on their own in interviews.

Remember this

Grown-up words: gluing ranges is merging intervals, the running totals are prefix sums, finding the first total that reaches K is a lower bound binary search, and the bag without repeats is the union of the ranges.