Question: Lift-to-stairs switch with an energy budget
Companies: Uber
Difficulty: Medium
Problem statement
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
e1energy and takest1time. Lift time isk * t1. -
Stairs: a floor climbed while your current energy is
Etakesceil(c / E)time, and then your energy drops bye2. Stair time is the sum over then - kstair floors. - Your energy may never drop below 0. A switch floor
kis valid only if you finish the stairs with energy>= 0, that isk * 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 costceil(4/2) = 2(energy 1), thenceil(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.
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
e1energy and costst1time. Afterkfloors you havek * e1energy and the lift time isk * t1. -
Stairs. A stair floor climbed with current energy
Ecostsceil(c / E)time, thenEdrops bye2. -
Energy never goes below 0. The switch floor
kis valid only ifk * e1 - (n - k) * e2 >= 0. If that holds, energy before every stair floor is at leaste2 >= 1, so the division is always defined. -
Objective. Return the smallest
|lift_time - stair_time|over all validk. (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 = 0is never valid whenn >= 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)(twochain_sumcalls, each at most one jump). -
Total:
O(n + c)time,O(c)extra space. Withn, c <= 1e5this is about2e5operations. -
Brute force:
O(n^2)time,O(1)space. - Python ints do not overflow. In Java or C++,
k * e1andk * t1reach1e10, 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, stairsceil(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 = 1gives lift 1, stairs 0). -
Only
k = nis valid (e1tiny,e2huge), e.g.n=3, e1=1, t1=5, e2=100, c=10->15. -
All stair floors cost 1 (
k * e1far abovec): the jump branch ofchain_sumis 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, stairsceil(1/1) = 1, answer0. -
Randomized check against
brute_forcefor small values (in the code).
Follow-ups
-
Minimize total time instead of the difference. Keep everything and change the score to
k * t1 + stairs. TheO(n + c)bound does not change. Mention that total time is not monotone ink, so you still need to scan everyk(no binary search). -
Report the switch floor too, or all floors that tie. Track
argminalongsidebest. For ties, pick the smallestk(least lift use) or ask the interviewer. -
Energy must stay strictly positive, not just non-negative. Change the validity test to
bottom > 0. A variant only requires energy> 0before each stair floor: then the test isbottom + e2 > 0, and the chain sum still works because every term has a positive argument. -
Multiple switches (lift, stairs, lift, ...). The closed form breaks. Use DP over
(floor, energy), with energy capped atcbecause energy abovecbehaves likecfor cost. That isO(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. -
Many queries with different
nand the samee1, t1, e2, c. Thechaintable depends only oncande2, so build it once. Each query is thenO(n). If queries also sharenranges, precompute answers pernoffline.
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 >= ccosts 1) without allocating a1e10array. - Tests edge cases: only
k = nvalid,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.