Max purchases from a sorted price row

Question: Max purchases from a sorted price row

Companies: Uber

Difficulty: Medium

Problem statement

A market hall has n stalls in a row, numbered 1..n. Stall i sells a single kind of item at price prices[i-1], and the prices are non-decreasing from left to right. A shopper may buy at most one item per stall.

You are given q queries as two arrays of length q, pos and amount. For query j, the shopper starts at stall pos[j] and may visit only stalls pos[j], pos[j]+1, ..., n. Return an array where entry j is the maximum number of items the shopper can buy in query j with total cost at most amount[j].

Implement:

def max_items(prices: list[int], pos: list[int], amount: list[int]) -> list[int]: ...

Constraints

• 1 <= n, q <= 10^5
• 0 <= prices[i] <= 10^9, and prices[i] <= prices[i+1]
• 1 <= pos[j] <= n
• 0 <= amount[j] <= 10^14

Example 1

prices = [2, 3, 5, 8]
pos    = [1, 2, 4]
amount = [10, 7, 7]
output = [3, 1, 0]

From stall 1 with 10: buy 2 + 3 + 5 = 10, so 3 items. From stall 2 with 7: 3 fits, 3 + 5 does not, so 1. From stall 4 with 7: the only item costs 8, so 0.

Example 2

prices = [1, 1, 1]
pos    = [1, 3]
amount = [100, 0]
output = [3, 0]

A brute-force walk per query is O(n * q). Aim for better.

Approach

Key observation. Prices never go down along the row. From a starting stall, the cheapest item you can still buy is the one at the starting stall, the next cheapest is the one after it, and so on. So the best plan for any budget is always "buy at the start, then the next stall, then the next", stopping when the next price no longer fits. The answer to a query is the largest k such that the first k prices of the suffix fit in the budget.

Why greedy is safe: any set of k stalls from the suffix costs at least as much as the first k stalls of the suffix, because the i-th cheapest chosen price is at least the i-th price of the suffix. If any k items fit, the first k fit.

Brute force. For each query, walk from pos and add prices until the budget runs out. That is O(n) per query, O(n * q) total. With n and q both up to 10^5 this is 10^10 steps, too slow, but it is the natural first answer and a good test oracle.

Optimal. Build prefix sums once: pre[i] = prices[0] + ... + prices[i-1]. The cost of buying stalls s .. s+k-1 (0-indexed s = pos - 1) is pre[s + k] - pre[s]. That is non-decreasing in k because prices are non-negative, so binary search for the largest end index e with pre[e] <= pre[s] + amount. The answer is e - s. In Python that is one bisect_right call per query.

Invariants.

• pre is non-decreasing (prices are non-negative), so bisect is valid.
• The search is restricted to e in [s, n], so the answer is never negative and never runs past the end.
• Python ints do not overflow; in Java or C++ use 64-bit sums (10^5 prices of 10^9 each exceed 32 bits).

Code

from bisect import bisect_right
from itertools import accumulate


def max_items(prices: list[int], pos: list[int], amount: list[int]) -> list[int]:
    """For each query j, the most items buyable from stalls pos[j]..n
    (1-indexed, at most one per stall) with total cost <= amount[j].
    prices must be non-decreasing and non-negative."""
    n = len(prices)
    pre = [0, *accumulate(prices)]          # pre[i] = sum(prices[:i])
    result: list[int] = []
    for p, budget in zip(pos, amount):
        s = p - 1                           # 0-indexed start
        if s < 0 or s >= n or budget < 0:
            result.append(0)
            continue
        # Largest e in [s, n] with pre[e] - pre[s] <= budget.
        e = bisect_right(pre, pre[s] + budget, lo=s, hi=n + 1) - 1
        result.append(e - s)
    return result


def max_items_brute(prices: list[int], pos: list[int], amount: list[int]) -> list[int]:
    out = []
    for p, budget in zip(pos, amount):
        count, spent = 0, 0
        for price in prices[p - 1:]:
            if spent + price > budget:
                break
            spent += price
            count += 1
        out.append(count)
    return out


if __name__ == "__main__":
    import random

    assert max_items([2, 3, 5, 8], [1, 2, 4], [10, 7, 7]) == [3, 1, 0]
    assert max_items([1, 1, 1], [1, 3], [100, 0]) == [3, 0]
    assert max_items([0, 0, 4], [1], [0]) == [2]
    assert max_items([5], [1, 1], [4, 5]) == [0, 1]
    for _ in range(500):
        n = random.randint(1, 30)
        prices = sorted(random.randint(0, 20) for _ in range(n))
        q = random.randint(1, 10)
        pos = [random.randint(1, n) for _ in range(q)]
        amount = [random.randint(0, 200) for _ in range(q)]
        assert max_items(prices, pos, amount) == max_items_brute(prices, pos, amount)
    print("ok")

Complexity

• Preprocessing (prefix sums): O(n) time, O(n) space.
• Each query: O(log n) time, O(1) extra space.
• Total: O(n + q log n) time, O(n + q) space including the output.
• Brute force for comparison: O(n * q) time, O(1) extra space.

Test cases

• prices = [2, 3, 5, 8], pos = [1, 2, 4], amount = [10, 7, 7] -> [3, 1, 0] (2+3+5 = 10 fits exactly; 3 fits but 3+5 does not; 8 > 7).
• prices = [1, 1, 1], pos = [1, 3], amount = [100, 0] -> [3, 0] (budget larger than everything buys the whole suffix; zero budget buys nothing).
• prices = [0, 0, 4], pos = [1], amount = [0] -> [2] (free items are always taken).
• prices = [5], pos = [1, 1], amount = [4, 5] -> [0, 1] (single stall; exact budget counts).
• Start at the last stall: prices = [1, 2, 3], pos = [3], amount = [3] -> [1].
• Large values: n = 10^5, every price 10^9, amount = 10^14 -> whole suffix; check 64-bit sums in other languages.
• Randomised check against the brute force on small sorted arrays.

Follow-ups

1. Prices are not sorted. Greedy still says "buy the k cheapest items in the suffix", but they are no longer the first k. Answer queries offline: sort queries by pos descending, insert prices into a Fenwick tree indexed by price rank (count and sum per rank) as pos moves left, then walk down the Fenwick tree (binary lifting) to find how many of the cheapest items fit. Handle a partial rank bucket with duplicate prices by taking remaining // price of them. O((n + q) log n). Online, use a persistent segment tree over price ranks.
2. Range queries [l, r] instead of a suffix. With sorted prices this is the same binary search with hi = r + 1. With unsorted prices, use a persistent segment tree (version r minus version l - 1) and descend by sum, O(log n) per query.
3. Price updates between queries (sorted order kept). Replace the prefix array with a Fenwick tree of prices. Find the answer with Fenwick binary lifting on target prefix(s) + amount. O(log n) per update and query.
4. Each stall has a stock count c[i]. Prefix sums of c[i] * prices[i] and of c[i]; binary search the last stall you can fully buy out, then add remaining // prices[e] items from the next stall (capped by its stock).
5. Maximise items, then minimise spend, or report which stalls. The greedy set is already the cheapest set of that size, so the spend is pre[e] - pre[s] and the stalls are s .. e-1; return them alongside the count.

Interviewer expectations

Mid-level (L4)

• Spots that sorted prices make the greedy "take the next stall" correct and can say why.
• Writes the O(n * q) loop, then moves to prefix sums with a binary search.
• Gets 1-indexing and the "exact budget counts" boundary right; tests empty-result and full-suffix cases.

Senior (L5)

• Goes straight to O(n + q log n) with a short exchange argument for the greedy.
• Uses bisect (or a hand-written search) with correct bounds so results never go negative or past n; mentions 64-bit overflow in typed languages.
• Handles follow-ups 3 and 4 (Fenwick for updates, stock counts) cleanly.

Staff (L6+)

• Solves the unsorted-prices version: offline Fenwick over price ranks with binary lifting, or a persistent segment tree for online and range queries.
• Discusses trade-offs between offline and online answering, memory of persistent structures, and duplicate prices in the descent.
• Validates with a brute-force oracle and randomised tests, and states the complexity of each variant precisely.