Best starting layer for a falling ball with an energy budget

Fix a start layer s. The ball has no choices after that: it passes layer
s, then s+1, and so on, as long as the energy left covers the next
layer's cost. So every start defines one forced window [s, end(s)), the
longest run with sum(costs[s:end]) <= energy, and the score is
sum(points[s:end]). The answer is the best score over all starts.

Brute force. Simulate from every start: O(n^2) time, O(1) space. This is
the right thing to say first, and it is also the oracle for testing.

Two pointers. Costs are non-negative, so if the ball starting at s gets
through layer e, the ball starting at s+1 does too: it skipped costs[s]
and has at least as much energy at every later layer. So end(s) never
moves left as s grows. Keep a sliding window with its running cost and
running points:

1. Extend end while spent + costs[end] <= energy.
2. Record gained as the score for this start.
3. Drop layer s from the window before moving to s+1.

If the ball cannot pass even its first layer (costs[s] > energy), the
window is empty, its score is 0, and the next start resets the window.

Invariants: spent == sum(costs[start:end]), gained == sum(points[start:end]),
spent <= energy, and end is the first layer the ball cannot pass.

Notes worth saying out loud:

• Points can be negative and the ball cannot stop early, so the answer can be
  negative (every layer passable and every point negative).
• Points do not refill energy here. If they did, it becomes the Gas Station
  problem (see follow-ups).
• The follow-up "the player may stop after any layer" is a sliding-window
  maximum over prefix sums of points, done with a monotonic deque, still O(n).

from collections import deque


def best_drop(costs: list[int], points: list[int], energy: int) -> tuple[int, int]:
    """Return (max_points, start_layer). The ball must keep falling while it can.

    start_layer is the smallest index achieving the maximum.
    """
    n = len(costs)
    if n != len(points):
        raise ValueError("costs and points must have the same length")
    best, best_start = 0 if n == 0 else None, 0
    end = 0              # first layer the ball does NOT pass when starting at `start`
    spent = 0            # sum(costs[start:end])
    gained = 0           # sum(points[start:end])
    for start in range(n):
        if end < start:  # window collapsed: nothing passed from the previous start
            end, spent, gained = start, 0, 0
        while end < n and spent + costs[end] <= energy:
            spent += costs[end]
            gained += points[end]
            end += 1
        if best is None or gained > best:
            best, best_start = gained, start
        if end > start:  # slide the left edge forward
            spent -= costs[start]
            gained -= points[start]
    return best, best_start


def best_drop_can_stop(costs: list[int], points: list[int], energy: int) -> int:
    """Follow-up: the player may stop the ball after any layer (or pass none)."""
    n = len(costs)
    prefix = [0] * (n + 1)
    for i, p in enumerate(points):
        prefix[i + 1] = prefix[i] + p
    best = 0
    window: deque[int] = deque()  # indices j into prefix, prefix[j] decreasing
    end = 0                       # prefix indices start+1 .. end are reachable stops
    spent = 0
    for start in range(n):
        if end < start:
            end, spent = start, 0
        while end < n and spent + costs[end] <= energy:
            spent += costs[end]
            end += 1
            while window and prefix[window[-1]] <= prefix[end]:
                window.pop()
            window.append(end)
        while window and window[0] <= start:
            window.popleft()
        if window:
            best = max(best, prefix[window[0]] - prefix[start])
        if end > start:
            spent -= costs[start]
    return best



if __name__ == "__main__":
    assert best_drop([2, 3, 1, 4], [5, -2, 4, 3], 6) == (7, 0)
    assert best_drop([10, 10], [1, 2], 5) == (0, 0)
    assert best_drop([1, 1, 1], [-1, -1, -1], 1) == (-1, 0)
    assert best_drop_can_stop([1, 1, 1, 1, 1], [4, -5, 3, 3, -9], 3) == 6
    assert best_drop([1, 1, 1, 1, 1], [4, -5, 3, 3, -9], 3) == (2, 0)
    print("ok")

• best_drop: O(n) time, since start and end each move forward at most
  n times. O(1) extra space.
• best_drop_can_stop: O(n) time, since each prefix index enters and leaves
  the deque once. O(n) space for the prefix sums and the deque.
• Brute force for comparison: O(n^2) time, O(1) space.

• costs=[2,3,1,4], points=[5,-2,4,3], energy=6 gives (7, 0): start 0 passes
  layers 0..2 (cost 6, points 7); start 2 also scores 7, the smaller start wins.
• costs=[10,10], points=[1,2], energy=5 gives (0, 0): no layer is passable.
• costs=[], points=[], energy=5 gives (0, 0).
• costs=[1,1,1], points=[-1,-1,-1], energy=1 gives (-1, 0): the ball must
  pass one layer and lose a point.
• costs=[0,0,5], points=[3,4,9], energy=0 gives (7, 0): zero-cost layers
  are always passed.
• Can stop: costs=[1,1,1,1,1], points=[4,-5,3,3,-9], energy=3 gives 6
  (start at layer 2, stop after layer 3). The forced version on the same input
  gives (2, 0), because a ball started at layer 2 must also pass layer 4.
• Can stop: all points negative gives 0 (stop before the first layer).
• Randomized check of both functions against the brute force on small arrays.

• Points refill energy (the Gas Station link). If passing layer i adds
  points[i] to energy, the window is no longer monotonic in general. For a
  circular tower that asks "from which start can the ball pass every layer",
  use the Gas Station greedy: one pass, reset the start whenever the running
  balance energy + sum(points - costs) goes negative, and check the total.
• Circular tower. The ball wraps from the last layer to the first. Double
  the arrays (or index mod n), cap each window at n layers, and run the same
  two pointers over starts 0..n-1. Still O(n).
• Stop any time. The player may stop after any layer. Score for a start is
  max(prefix[k] - prefix[s]) for k in [s, end(s)], a sliding-window max
  of prefix sums with a monotonic deque (best_drop_can_stop).
• Many queries for different initialEnergy. Precompute prefix costs. For
  one query, end(s) is a binary search on prefix costs, so O(n log n) per
  query without the window. For many queries, sort them and reuse the
  two-pointer pass per query, or precompute per-start answers for the distinct
  breakpoints if Q is large.
• Negative costs (layers that give energy). Monotonicity breaks, so two
  pointers are wrong. Simulate per start with an early exit, or use prefix sums
  plus a sparse table to find the first prefix that exceeds the budget by
  binary search on a range-maximum, O(n log n).

Mid-level (L4)

• Restates the rules: forced fall, any start, stop when energy is short.
• Gets a correct O(n^2) simulation, then moves to two pointers with a hint.
• Handles the empty window and negative totals.

Senior (L5)

• Argues why end(s) is monotonic (non-negative costs) before coding.
• Writes the O(n) window without a hint, with clear invariants and tie rules.
• Tests against a brute force and covers zero-cost and unpassable layers.
• Solves "stop any time" with the monotonic deque.

Staff (L6+)

• Separates which assumptions make the window valid (non-negative costs, no
  refill) and says what changes when each one breaks.
• Links the refill version to Gas Station and gives the greedy proof sketch.
• Discusses the multi-query version and picks an approach from the sizes of n
  and Q.