~/problems / Binary search / Prefix sums + binary search

How many items fit the budget

easy ~15 min Uber

A street market has n stalls in a line, numbered 1..n. Stall i sells a single item for prices[i - 1], and prices never decrease along the street. A shopper enters at some stall, walks to the end of the street, and may buy at most one item per stall, spending no more than their budget.

Implement max_items(prices: list[int], pos: list[int], amount: list[int]) -> list[int]. Query j is a shopper who enters at stall pos[j] (1-indexed) with budget amount[j]; they may buy from stalls pos[j]..n only. Return, for each query, the largest number of items they can buy.

  • 1 <= prices[i] <= 10**9, sorted non-decreasing.
  • 1 <= pos[j] <= n, 0 <= amount[j] <= 10**15.
prices = [2, 3, 3, 7, 10]
max_items(prices, [1, 2, 4, 5, 3], [8, 5, 100, 9, 0])
# [3, 1, 2, 0, 0]
# stall 1, budget 8:   2 + 3 + 3 = 8
# stall 2, budget 5:   3 (3 + 3 = 6 is over)
# stall 4, budget 100: 7 + 10
# stall 5, budget 9:   the only item costs 10

n and the number of queries go up to 100,000 each, so walking the stalls for every query is too slow.

Show hint

Because prices never decrease, the best buy from stall p is always the next few stalls in a row, p, p+1, .... With a prefix-sum array pre, you want the largest k with pre[p - 1 + k] - pre[p - 1] <= budget: binary search (bisect_right) for budget + pre[p - 1].

Topic: Prefix sums + binary search. Weighted sampling: bisect into cumulative weights.

Read the visual guide
0:00
Ctrl ' run · Ctrl ↵ submit
esc