Implement max profit with K transactions (DP)
Company: Citadel
Role: Data Scientist
Category: Coding & Algorithms
Difficulty: Medium
Interview Round: Technical Screen
Quick Answer: This question evaluates dynamic programming, algorithmic optimization, and complexity-analysis skills by requiring a bottom-up DP implementation that maximizes stock trading profit with at most k transactions while returning optimal buy-sell pairs under tie-breaking constraints.
Constraints
- 0 <= len(prices) <= 10000
- 0 <= k <= 100
- 0 <= prices[i] <= 10^9
- Transactions must satisfy buy_day < sell_day, and different transactions must be strictly disjoint in time.
Examples
Input: ([3, 2, 6, 5, 0, 3], 2)
Expected Output: (7, [(1, 2), (4, 5)])
Explanation: Buy on day 1 and sell on day 2 for profit 4, then buy on day 4 and sell on day 5 for profit 3. Total profit is 7.
Input: ([1, 3, 2, 0, 2], 2)
Expected Output: (4, [(0, 1), (3, 4)])
Explanation: The best two disjoint transactions are (0,1) for profit 2 and (3,4) for profit 2, totaling 4.
Input: ([1, 1, 2, 2], 1)
Expected Output: (1, [(0, 2)])
Explanation: Maximum profit is 1. Several trades achieve it, but (0,2) is lexicographically smaller than (0,3), (1,2), and (1,3).
Input: ([5, 4, 3, 2], 3)
Expected Output: (0, [])
Explanation: Prices only decrease, so the best profit is 0 and the empty transaction list is lexicographically smallest.
Input: ([], 3)
Expected Output: (0, [])
Explanation: With no days, no transaction is possible.
Input: ([5], 4)
Expected Output: (0, [])
Explanation: With fewer than 2 days, no transaction is possible.
Input: ([2, 1, 2, 0, 1], 0)
Expected Output: (0, [])
Explanation: When k = 0, no transactions are allowed.
Input: ([1, 2, 1, 2], 10)
Expected Output: (2, [(0, 1), (2, 3)])
Explanation: Even with large k, at most two profitable disjoint transactions exist here: (0,1) and (2,3).
Solution
def solution(prices, k):
n = len(prices)
if k <= 0 or n < 2:
return (0, [])
k = min(k, n // 2)
if k == 0:
return (0, [])
# dp[t][i] = maximum profit from days i..n-1 using at most t transactions.
dp = [[0] * (n + 1) for _ in range(k + 1)]
# take[t][i] tells whether the lexicographically smallest optimal solution
# for state (t, i) starts by buying on day i.
take = [[False] * n for _ in range(k + 1)]
# If take[t][i] is True, sell_choice[t][i] is the chosen sell day.
sell_choice = [[-1] * n for _ in range(k + 1)]
for t in range(1, k + 1):
best_value = prices[n - 1] + dp[t - 1][n]
best_sell = n - 1
dp[t][n] = 0
dp[t][n - 1] = 0
for i in range(n - 2, -1, -1):
candidate = best_value - prices[i]
skip = dp[t][i + 1]
# Tie-breaking:
# - If candidate profit is larger, take it.
# - If profits tie and are positive, taking at day i is lexicographically
# smaller than skipping to a later buy day.
# - If both are 0, prefer skipping so we return the shorter list [].
if candidate > skip or (candidate == skip and candidate > 0):
dp[t][i] = candidate
take[t][i] = True
sell_choice[t][i] = best_sell
else:
dp[t][i] = skip
take[t][i] = False
# Update the best future sell option for the next iteration.
value = prices[i] + dp[t - 1][i + 1]
if value > best_value or (value == best_value and i < best_sell):
best_value = value
best_sell = i
transactions = []
t = k
i = 0
while t > 0 and i < n - 1 and dp[t][i] > 0:
if take[t][i]:
s = sell_choice[t][i]
transactions.append((i, s))
i = s + 1
t -= 1
else:
i += 1
return (dp[k][0], transactions)
Time complexity: O(nk). Space complexity: O(nk).
Hints
- Use exact-transaction DP states: cash[t] for the best profit after exactly t completed transactions and no stock in hand, and hold[t] for the best balance after opening transaction t.
- The O(n^2 k) version tries every earlier buy day for every sell day. You can avoid that by carrying forward the best value of cash[t-1] - prices[i] as you sweep from left to right.