State-machine DP: buy, sell, cool down

import random

from trades import max_profit_k

PRICES = [5, 3, 6, 8, 2, 7, 4, 9]


def every_plan(prices, k, fee):
    """The answer the tests trust for a few days: every plan of waits, buys and sells, with at most k sales."""

    def day(d, holding, left):
        if d == len(prices):
            return float("-inf") if holding else 0
        best = day(d + 1, holding, left)
        if holding:
            best = max(best, prices[d] - fee + day(d + 1, False, left))
        elif left > 0:
            best = max(best, -prices[d] + day(d + 1, True, left - 1))
        return best

    return day(0, False, k)


def synthetic(n, seed):
    """Made-up prices from a small linear congruential generator."""
    prices, x = [], seed
    for _ in range(n):
        x = (x * 1_103_515_245 + 12_345) % 2_147_483_648
        prices.append(1 + x % 10_000)
    return prices


def test_the_examples():
    """the examples of the prompt, and more sales allowed on the same prices"""
    assert max_profit_k(PRICES, 0, 0) == 0
    assert max_profit_k(PRICES, 1, 0) == 7
    assert max_profit_k(PRICES, 2, 1) == 10
    assert max_profit_k([9, 7, 4, 2], 2, 0) == 0
    assert [max_profit_k(PRICES, k, 0) for k in range(5)] == [0, 7, 12, 15, 15]
    assert [max_profit_k(PRICES, k, 1) for k in range(5)] == [0, 6, 10, 12, 12]


def test_one_day_and_a_large_fee():
    """one day allows no sale; a fee larger than every rise makes trading lose"""
    assert max_profit_k([5], 3, 0) == 0
    assert max_profit_k(PRICES, 3, 7) == 0


def test_small_cases_against_every_plan():
    """200 random cases of up to 10 days, against every plan"""
    rng = random.Random(53)
    for _ in range(200):
        prices = [rng.randint(1, 15) for _ in range(rng.randint(1, 10))]
        k, fee = rng.randint(0, 5), rng.randint(0, 3)
        assert max_profit_k(prices, k, fee) == every_plan(prices, k, fee), f"max_profit_k({prices}, {k}, {fee})"


def test_two_thousand_days_with_ten_sales():
    """2,000 made-up prices, at most 10 sales and a fee of 30"""
    assert max_profit_k(synthetic(2_000, 7), 10, 30) == 99_159


def test_when_k_stops_mattering():
    """20,000 days with k = 20,000: a table of n × k entries would need 4 × 10^8 steps"""
    prices = synthetic(20_000, 99)
    assert max_profit_k(prices, 20_000, 40) == 32_691_476
    assert max_profit_k(prices, 3, 40) == 29_876

import { test, assert } from 'mysmartcopilot:test';
import { maxProfitK } from './trades.mjs';

const PRICES = [5, 3, 6, 8, 2, 7, 4, 9];

/** The answer the tests trust for a few days: every plan of waits, buys and sells, with at most k sales. */
function everyPlan(prices, k, fee) {
  function day(d, holding, left) {
    if (d === prices.length) return holding ? -Infinity : 0;
    let best = day(d + 1, holding, left);
    if (holding) best = Math.max(best, prices[d] - fee + day(d + 1, false, left));
    else if (left > 0) best = Math.max(best, -prices[d] + day(d + 1, true, left - 1));
    return best;
  }
  return day(0, false, k);
}

/** Made-up prices from a small linear congruential generator (BigInt: the product passes 2 ** 53). */
function synthetic(n, seed) {
  const prices = [];
  let x = BigInt(seed);
  for (let i = 0; i < n; i++) {
    x = (x * 1_103_515_245n + 12_345n) % 2_147_483_648n;
    prices.push(1 + Number(x % 10_000n));
  }
  return prices;
}

test('the examples of the prompt, and more sales allowed on the same prices', () => {
  assert.equal(maxProfitK(PRICES, 0, 0), 0);
  assert.equal(maxProfitK(PRICES, 1, 0), 7);
  assert.equal(maxProfitK(PRICES, 2, 1), 10);
  assert.equal(maxProfitK([9, 7, 4, 2], 2, 0), 0);
  assert.deepEqual([0, 1, 2, 3, 4].map((k) => maxProfitK(PRICES, k, 0)), [0, 7, 12, 15, 15]);
  assert.deepEqual([0, 1, 2, 3, 4].map((k) => maxProfitK(PRICES, k, 1)), [0, 6, 10, 12, 12]);
});

test('one day allows no sale; a fee larger than every rise makes trading lose', () => {
  assert.equal(maxProfitK([5], 3, 0), 0);
  assert.equal(maxProfitK(PRICES, 3, 7), 0);
});

test('150 random cases of up to 9 days, against every plan', () => {
  let s = 53;
  const next = (n) => {
    s = (s * 48_271) % 2_147_483_647;
    return s % n;
  };
  for (let t = 0; t < 150; t++) {
    const prices = Array.from({ length: 1 + next(9) }, () => 1 + next(15));
    const k = next(6);
    const fee = next(4);
    assert.equal(maxProfitK(prices, k, fee), everyPlan(prices, k, fee), `maxProfitK([${prices}], ${k}, ${fee})`);
  }
});

test('2,000 made-up prices, at most 10 sales and a fee of 30', () => {
  assert.equal(maxProfitK(synthetic(2_000, 7), 10, 30), 99_159);
});

test('20,000 days with k = 20,000: a table of n × k entries would need 4 × 10^8 steps', () => {
  const prices = synthetic(20_000, 99);
  assert.equal(maxProfitK(prices, 20_000, 40), 32_691_476);
  assert.equal(maxProfitK(prices, 3, 40), 29_876);
});