Lift-to-stairs switch with an energy budget

0
⑂ RemixPractice

You start on floor 0 of a building with 0 energy and need to reach floor n. You ride the lift from floor 0 up to some floor k (0 <= k <= n), then walk the stairs from floor k to floor n.

  • Lift: each floor adds e1 energy and takes t1 time. Lift time is k * t1.
  • Stairs: a floor climbed while your current energy is E takes ceil(c / E) time, and then your energy drops by e2. Stair time is the sum over the n - k stair floors.
  • Your energy may never drop below 0. A switch floor k is valid only if you finish the stairs with energy >= 0, that is k * e1 - (n - k) * e2 >= 0.

Implement:

def min_time_difference(n: int, e1: int, t1: int, e2: int, c: int) -> int:
    ...

Return the minimum value of |lift time - stair time| over all valid switch floors k. Riding the lift all the way (k = n, no stairs) is always valid.

Constraints: 1 <= n, e1, t1, e2, c <= 10^5. An O(n^2) simulation is too slow for the largest inputs.

Example 1

n = 3, e1 = 2, t1 = 5, e2 = 1, c = 4   ->   1
  • k = 0: invalid (no energy).
  • k = 1: lift time 5, energy 2. Stairs cost ceil(4/2) = 2 (energy 1), then ceil(4/1) = 4 (energy 0). Stair time 6, difference 1.
  • k = 2: lift time 10, energy 4, one stair floor costs 1. Difference 9.
  • k = 3: lift time 15, no stairs. Difference 15.

Example 2

n = 4, e1 = 3, t1 = 2, e2 = 2, c = 7   ->   0

Switch at k = 2: lift time 4, energy 6. Stairs cost ceil(7/6) = 2 (energy 4), then ceil(7/4) = 2 (energy 2). Stair time 4, difference 0. k = 1 is invalid: energy 3 cannot cover three stair floors at 2 each.

UberOnline assessmentMediumLast asked Feb 2026

What L5 looks like

  • Sees that the stair energies are an arithmetic sequence and turns the stair time into a prefix-sum difference.
  • Handles the large-energy case (E >= c costs 1) without allocating a 1e10 array.
  • Tests edge cases: only k = n valid, bottom == 0, all-ones stair cost. Mentions 64-bit overflow in other languages.
Interviewer expectations for every level ↓

Approach

Clarify first. You start on floor 0 with 0 energy and must end on floor n. You ride the lift for floors 0 -> k, then take the stairs for floors k -> n. Pin down these points:

  • Lift. Each lift floor adds e1 energy and costs t1 time. After k floors you have k * e1 energy and the lift time is k * t1.
  • Stairs. A stair floor climbed with current energy E costs ceil(c / E) time, then E drops by e2.
  • Energy never goes below 0. The switch floor k is valid only if k * e1 - (n - k) * e2 >= 0. If that holds, energy before every stair floor is at least e2 >= 1, so the division is always defined.
  • Objective. Return the smallest |lift_time - stair_time| over all valid k. (Some versions ask for the smallest total time instead. See Follow-ups. The same machinery answers both.)
  • k = n (lift all the way, no stairs) is always valid, so an answer always exists. k = 0 is never valid when n >= 1.

Check with the reported example: n = 5, e1 = 2, t1 = 3, e2 = 4, c = 5. k = 3 is invalid (6 - 2*4 < 0). k = 4 gives lift 12 and stairs ceil(5/8) = 1, a difference of 11. k = 5 gives 15. The answer is 11.

Brute force. Try every k, check validity, then simulate the n - k stair floors. That is O(n^2) in the worst case. With n up to 1e5 that is about 5e9 steps, too slow. Start here, then say where the time goes: the stair sum is recomputed from scratch for every k.

Key observation 1: the stair energies form an arithmetic sequence. For switch floor k, the energies on the stairs are top, top - e2, top - 2*e2, ..., with top = k * e1. After the last floor, energy is bottom = top - (n - k) * e2 >= 0. Define

chain(v) = ceil(c/v) + ceil(c/(v - e2)) + ceil(c/(v - 2*e2)) + ... (all terms with a positive argument).

Then stair_time(k) = chain(top) - chain(bottom), exactly like a prefix-sum difference. chain(v) = ceil(c/v) + chain(v - e2), so within one residue class mod e2 it is a running sum.

Key observation 2: above c, every stair floor costs exactly 1. top can be as large as n * e1 = 1e10, too big to tabulate. But ceil(c/E) = 1 whenever E >= c. So we only tabulate chain(v) for v < c. For v >= c, count how many terms in the chain are still >= c (t = (v - c) // e2 + 1), add t, and jump to chain(v - t * e2), which is below c and in the table.

Result. Build the table in O(c). Then each k is O(1): two chain lookups and one subtraction. The total is O(n + c).

Invariant. chain[v] for 0 <= v < c equals the exact stair time of starting a walk at energy v and going down until energy reaches <= 0. A walk from top that stops at bottom is the walk from top minus the walk from bottom, because both walks visit the same residue class and the walk from bottom is the tail of the walk from top.

Code

def min_time_difference(n: int, e1: int, t1: int, e2: int, c: int) -> int:
    """Smallest |lift time - stair time| over every valid switch floor k."""
    # chain[v] = ceil(c/v) + ceil(c/(v-e2)) + ... while the argument is > 0.
    # Tabulate only v < c: at energy >= c every stair floor costs exactly 1.
    chain = [0] * c
    for v in range(1, c):
        prev = chain[v - e2] if v - e2 > 0 else 0
        chain[v] = (c + v - 1) // v + prev

    def chain_sum(v: int) -> int:
        if v <= 0:
            return 0
        if v < c:
            return chain[v]
        floors_at_one = (v - c) // e2 + 1  # terms v, v-e2, ... that are >= c
        return floors_at_one + chain_sum(v - floors_at_one * e2)  # now v < c

    best: int | None = None
    for k in range(n + 1):  # ride the lift for floors 0..k
        top = k * e1  # energy when stepping off the lift
        bottom = top - (n - k) * e2  # energy after the last stair floor
        if bottom < 0:  # energy would drop below zero
            continue
        stairs = chain_sum(top) - chain_sum(bottom)
        diff = abs(k * t1 - stairs)
        if best is None or diff < best:
            best = diff
    assert best is not None  # k = n is always valid
    return best


def brute_force(n: int, e1: int, t1: int, e2: int, c: int) -> int:
    """O(n^2) reference: simulate the stairs for every switch floor."""
    best: int | None = None
    for k in range(n + 1):
        energy, stairs = k * e1, 0
        for _ in range(n - k):
            stairs += -(-c // energy) if energy > 0 else 0
            energy -= e2
        if energy >= 0:
            diff = abs(k * t1 - stairs)
            best = diff if best is None else min(best, diff)
    assert best is not None
    return best


if __name__ == "__main__":
    import random

    assert min_time_difference(5, 2, 3, 4, 5) == 11
    assert min_time_difference(4, 3, 2, 2, 7) == 0
    assert min_time_difference(1, 1, 1, 1, 1) == 1
    for _ in range(2000):
        args = [random.randint(1, 12) for _ in range(5)]
        assert min_time_difference(*args) == brute_force(*args), args
    min_time_difference(10**5, 10**5, 10**5, 1, 10**5)  # runs well under 1s
    print("all tests passed")

chain_sum recurses at most once: after the jump, v < c.

Complexity

  • Table build: O(c) time, O(c) space.
  • Per switch floor: O(1) (two chain_sum calls, each at most one jump).
  • Total: O(n + c) time, O(c) extra space. With n, c <= 1e5 this is about 2e5 operations.
  • Brute force: O(n^2) time, O(1) space.
  • Python ints do not overflow. In Java or C++, k * e1 and k * t1 reach 1e10, so use 64-bit integers.

Test cases

  • n=5, e1=2, t1=3, e2=4, c=5 -> 11 (switch at 4: lift 12, stairs 1).
  • n=4, e1=3, t1=2, e2=2, c=7 -> 0 (switch at 2: lift 4, stairs ceil(7/6) + ceil(7/4) = 2 + 2 = 4).
  • n=1, e1=1, t1=1, e2=1, c=1 -> 1 (switch at 0 is invalid; k = 1 gives lift 1, stairs 0).
  • Only k = n is valid (e1 tiny, e2 huge), e.g. n=3, e1=1, t1=5, e2=100, c=10 -> 15.
  • All stair floors cost 1 (k * e1 far above c): the jump branch of chain_sum is used, e.g. n=10^5, e1=10^5, t1=10^5, e2=1, c=10^5 -> 100002.
  • Energy ends exactly at 0 (bottom == 0) is allowed: n=2, e1=1, t1=1, e2=1, c=1 -> switch at 1, lift 1, stairs ceil(1/1) = 1, answer 0.
  • Randomized check against brute_force for small values (in the code).

Follow-ups

  1. Minimize total time instead of the difference. Keep everything and change the score to k * t1 + stairs. The O(n + c) bound does not change. Mention that total time is not monotone in k, so you still need to scan every k (no binary search).
  2. Report the switch floor too, or all floors that tie. Track argmin alongside best. For ties, pick the smallest k (least lift use) or ask the interviewer.
  3. Energy must stay strictly positive, not just non-negative. Change the validity test to bottom > 0. A variant only requires energy > 0 before each stair floor: then the test is bottom + e2 > 0, and the chain sum still works because every term has a positive argument.
  4. Multiple switches (lift, stairs, lift, ...). The closed form breaks. Use DP over (floor, energy), with energy capped at c because energy above c behaves like c for cost. That is O(n * c) states, so ask about limits first. If only the total time matters and lift time per floor is constant, show that one switch is optimal when stairs never add energy.
  5. Many queries with different n and the same e1, t1, e2, c. The chain table depends only on c and e2, so build it once. Each query is then O(n). If queries also share n ranges, precompute answers per n offline.

Interviewer expectations

Mid-level (L4)

  • Restates the rules correctly, including energy starting at 0 and the validity check.
  • Writes the O(n^2) simulation cleanly and gets the reported example right.
  • Uses integer ceiling division ((c + E - 1) // E), not floating point.

Senior (L5)

  • Sees that the stair energies are an arithmetic sequence and turns the stair time into a prefix-sum difference.
  • Handles the large-energy case (E >= c costs 1) without allocating a 1e10 array.
  • Tests edge cases: only k = n valid, bottom == 0, all-ones stair cost. Mentions 64-bit overflow in other languages.

Staff (L6+)

  • States the invariant behind chain(top) - chain(bottom) and why one jump is enough.
  • Discusses how the objective variants (difference, total time, strict energy) change only the score or the validity test.
  • Reasons about generalizations (multiple switches, many queries) and their state-space cost before coding them.

Make it yours

Write each chapter yourself, check it against this answer, then let the coach grill you.