stock_trade.py

import numpy as np


def k_best_trades(b, s, k):
    n = len(b)
    dp = np.zeros(shape=(k+1, n))
    buy_idx = np.full(shape=(k+1, n), fill_value=-1)
    best = np.full(shape=(k+1,), fill_value=np.inf)
    best_buy_idx = np.full(shape=(k+1,), fill_value=-1)

    for t in np.arange(1, k+1):
        best[t] = -b[0]
        best_buy_idx[t] = 0

    for i in np.arange(1, n):
        for t in np.arange(1, k+1):
            dp[t][i] = dp[t][i-1]

            cand = s[i] + best[t]
            if cand > dp[t][i]:
                dp[t][i] = cand
                buy_idx[t][i] = best_buy_idx[t]

            tmp = dp[t-1][i] - b[i]
            if tmp > best[t]:
                best[t] = tmp
                best_buy_idx[t] = i

    t, i, trades = k, n-1, list()
    while t and i:
        if dp[t][i] == dp[t][i-1]:
            i -= 1
        else:
            sell = i
            buy  = buy_idx[t][i]
            trades.append((buy, sell))
            i = buy - 1
            t -= 1

    trades.reverse()
    return dp[k][n-1], trades

b = [3,14,9,21,7,11,6,18,2,20,5,16]
s = [24,25,17,8,24,13,22,1,15,19,19,12]

profit, pairs = k_best_trades(b, s, k=2)