Max flight distance with limited energy: Frontier Interview's solution

Question: Max flight distance with limited energy

Companies: Google

Difficulty: Hard

Problem statement

You are planning a trip over n days. dist[i] is how far you can fly on day i (a non-negative integer). You start with k units of energy, and k is also the most energy you can hold.

Each day you do exactly one of two things:

Fly: spend 1 energy and cover dist[i]. You cannot fly with 0 energy.
Rest: gain 1 energy, but never more than k.

Implement max_fly_distance(dist: list[int], k: int) -> int, which returns the largest total distance you can cover over the n days.

Constraints: 0 <= n <= 2000, 0 <= k <= 10^9, 0 <= dist[i] <= 10^6.

Example 1

dist = [4, 1, 7, 3], k = 1
output: 11

Fly on day 0 (4), rest on day 1, fly on day 2 (7), rest on day 3.

Example 2

dist = [1, 1, 10, 10, 10], k = 1
output: 21

Fly on day 0 (1), rest, fly (10), rest, fly (10). Resting on day 0 would gain nothing, because you are already at full energy.

A follow-up may ask you to return the day-by-day plan as well.

Approach

Each day is a binary choice, fly or rest, and the only thing that links one
day to the next is how much energy you have. That makes energy the DP state.

Brute force. Try all 2^n fly/rest strings, simulate the energy, discard
any string that flies at 0 energy, keep the best total. O(2^n · n). Say it,
then move on.

DP over (day, energy). Let best[e] be the largest distance you can have
covered after the days processed so far, ending with exactly e energy
(-inf if unreachable). Start with best[k] = 0. For day i with distance
d, from every reachable e:

fly (needs e > 0): next[e - 1] = max(next[e - 1], best[e] + d)
rest: next[min(k, e + 1)] = max(..., best[e])

The answer is max(best) after the last day. Only the previous row is needed,
so space is O(k).

Why it is correct: the future depends only on the day index and the current
energy, not on which days you flew earlier, so two histories with the same
energy are interchangeable and we keep only the better one.

Cap the state at n. Energy above n can never be spent, so use
cap = min(k, n). That bounds the work at O(n · min(n, k)) even when k is
huge, and if k >= n the answer is just sum(dist) (fly every day).

A useful invariant (for discussion, and for the follow-ups). Because you
start full and every day changes energy by exactly ±1, a fly set is feasible
iff every contiguous window of L days has at most (L + k) / 2 fly days
(flies minus rests in any window is at most k). The cap at k is what makes
it a window condition rather than a prefix condition: hitting the cap
"forgets" earlier rests.

Distances are non-negative in the base problem; if they can be negative the
DP still works, because resting is always allowed.

Code

from typing import List, Tuple

NEG = float("-inf")


def max_fly_distance(dist: List[int], k: int) -> int:
    """Best total distance. Start with k energy; fly costs 1, rest restores 1 (cap k)."""
    if k <= 0 or not dist:
        return 0
    cap = min(k, len(dist))           # energy above n can never be spent
    # best[e] = max distance so far, ending the day with e energy
    best = [NEG] * (cap + 1)
    best[cap] = 0
    for d in dist:
        nxt = [NEG] * (cap + 1)
        for e in range(cap + 1):
            if best[e] == NEG:
                continue
            if e > 0:                                  # fly
                nxt[e - 1] = max(nxt[e - 1], best[e] + d)
            r = min(cap, e + 1)                        # rest
            nxt[r] = max(nxt[r], best[e])
        best = nxt
    return int(max(best))


def plan_schedule(dist: List[int], k: int) -> Tuple[int, List[str]]:
    """Same DP, keeping the full table so we can return one optimal plan."""
    n = len(dist)
    if k <= 0 or n == 0:
        return 0, ["rest"] * n
    cap = min(k, n)
    dp = [[NEG] * (cap + 1) for _ in range(n + 1)]
    dp[0][cap] = 0
    for i, d in enumerate(dist):
        for e in range(cap + 1):
            v = dp[i][e]
            if v == NEG:
                continue
            if e > 0 and v + d > dp[i + 1][e - 1]:
                dp[i + 1][e - 1] = v + d
            r = min(cap, e + 1)
            if v > dp[i + 1][r]:
                dp[i + 1][r] = v
    e = max(range(cap + 1), key=lambda x: dp[n][x])
    total = dp[n][e]
    plan: List[str] = []
    for i in range(n, 0, -1):                          # walk back
        d = dist[i - 1]
        if e + 1 <= cap and dp[i - 1][e + 1] != NEG and dp[i - 1][e + 1] + d == dp[i][e]:
            plan.append("fly")
            e += 1
            continue
        prevs = [e - 1] + ([cap] if e == cap else [])
        for p in prevs:
            if 0 <= p <= cap and dp[i - 1][p] == dp[i][e]:
                e = p
                break
        plan.append("rest")
    plan.reverse()
    return int(total), plan


if __name__ == "__main__":
    assert max_fly_distance([4, 1, 7, 3], 1) == 11
    assert max_fly_distance([3, 9, 2, 8, 1], 2) == 20
    assert max_fly_distance([1, 1, 10, 10, 10], 1) == 21
    assert max_fly_distance([5, 6, 7], 5) == 18
    assert max_fly_distance([], 3) == 0
    assert plan_schedule([4, 1, 7, 3], 1) == (11, ["fly", "rest", "fly", "rest"])
    print("ok")

Complexity

max_fly_distance: O(n · min(n, k)) time, O(min(n, k)) space.
plan_schedule: O(n · min(n, k)) time and space (the full table is kept so
the plan can be walked back).
Brute force, for comparison: O(2^n · n).

Test cases

dist = [4, 1, 7, 3], k = 1 → 11 (fly, rest, fly, rest).
dist = [3, 9, 2, 8, 1], k = 2 → 20 (fly 3 and 9, rest, fly 8).
dist = [1, 1, 10, 10, 10], k = 1 → 21 (fly day 1, rest, fly, rest, fly).
Flying the cheap first day is free because the cap would waste a rest anyway.
dist = [5, 6, 7], k = 5 → 18 (k ≥ n: fly every day).
dist = [5, 5, 5], k = 0 → 0 (never able to fly).
dist = [], k = 3 → 0.
dist = [0, 0, 9], k = 1 → 9.
Randomized check: 3,000 random cases with n ≤ 10, k ≤ 5 match the 2^n brute
force, and every returned plan replays to the same total without going
below zero energy.

Follow-ups

Return the schedule, not just the total. Keep the full n × (cap + 1)
table (or a parent pointer per cell) and walk back from the best final
energy, as plan_schedule does. Space becomes O(n · min(n, k)).
Resting restores r units, or flying day i costs c[i]. Same DP;
the transitions become e → min(k, e + r) and e → e - c[i] (needs
e >= c[i]). The state bound is still k.
Very large n and k (say 10^5 each). O(n · k) is too slow. Use the
window characterization: you are picking a fly set where every window has
flies − rests ≤ k. Without the cap this is the classic exchange greedy:
walk left to right, tentatively fly every day, and when energy would go
negative, turn the cheapest flown day so far (min-heap) into a rest, which
gains 2 energy. The cap breaks the pure prefix argument, because a rest
taken at full energy is wasted, so a candidate must explain how it handles
that (for example by splitting at days where energy is forced to k).
Expect to argue it, not prove it fully, and to validate it against the DP
on random inputs.
Start with s < k energy, or must end with at least m energy. Seed
best[s] = 0 instead of best[k]; take the answer as
max(best[m:]).
Weather days where flying is forbidden, or a second action (e.g. "glide"
that costs 0 energy but gives half distance). Add or remove transitions;
the state stays (day, energy).

Interviewer expectations

Mid-level (L4)

Simulates the rules correctly, including the cap at k and no flying at 0.
Gets from brute force to the (day, energy) DP with light prompting.
Writes working code with a rolling array and handles empty input and k = 0.

Senior (L5)

Names the state and explains why it is sufficient (the future depends only
on day and energy).
Caps energy at min(k, n) and notes k >= n means sum of all distances.
Reconstructs the schedule and tests against a brute force on small inputs.

Staff (L6+)

States the window invariant (flies − rests ≤ k in every window) and uses it
to reason about the greedy/heap approach for large n and k.
Generalizes cleanly to variable rest gain, variable fly cost, and end-energy
constraints without rewriting the core.
Discusses where the greedy needs proof and how to de-risk it (DP as an
oracle in randomized tests).