Best starting layer for a falling ball with an energy budget: Frontier Interview's solution

Question: Best starting layer for a falling ball with an energy budget

Companies: Google

Difficulty: Hard

Problem statement

A ball falls through a tower of n layers, numbered 0 (top) to n - 1 (bottom). You get two integer arrays of length n and one integer:

costs[i] (>= 0): the energy the ball spends to pass through layer i.
points[i] (may be negative): the score added when the ball passes through layer i.
energy (>= 0): the energy the ball starts with.

You drop the ball at any layer s you choose. From there it must keep falling: it passes layer s, then s + 1, and so on, as long as its remaining energy is at least the cost of the next layer. It stops at the first layer it cannot afford, or after the bottom layer. Points do not give energy back.

Implement best_drop(costs: list[int], points: list[int], energy: int) -> tuple[int, int] that returns the best total score over all starting layers and the smallest starting layer that achieves it. If the ball can pass no layer from any start, the score is 0. For an empty tower return (0, 0).

Constraints: 0 <= n <= 10^5, 0 <= costs[i] <= 10^4, -10^4 <= points[i] <= 10^4, 0 <= energy <= 10^9. Aim for better than O(n^2).

Example 1

costs  = [2, 3, 1, 4]
points = [5, -2, 4, 3]
energy = 6

Output: (7, 0). From layer 0 the ball passes layers 0, 1, 2 (cost 6) for 5 - 2 + 4 = 7. From layer 2 it passes layers 2 and 3 (cost 5) for 7 as well; layer 0 is the smaller start.

Example 2

costs  = [1, 1, 1]
points = [-1, -1, -1]
energy = 1

Output: (-1, 0). Every start passes exactly one layer, and the ball cannot choose to stop.

Likely follow-up: the player may stop the ball after any layer. Return the best score in that case.

Reconstructed from the reported question; details may differ from what interviewers ask.

Approach

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:

Extend end while spent + costs[end] <= energy.
Record gained as the score for this start.
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).

Code

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")

Complexity

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.

Test cases

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.

Follow-ups

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).

Interviewer expectations

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.