Best time to buy and sell stock using strategy

Array Sliding Window Prefix Sum

59.8% acceptance

Feb 25, 2026

351

You are given two integer arrays prices and strategy, where: prices[i] is the price of a given stock on the ith day. strategy[i] represents a trading action: -1 buy, 0 hold, 1 sell. You may perform at most one modification: Select exactly k consecutive elements. Set first k/2 to 0 (hold), last k/2 to 1 (sell). Profit = sum of strategy[i] * prices[i]. Return the maximum possible profit.

Solution

Rust

Time O(n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn max_profit(prices: Vec<i32>, strategy: Vec<i32>, k: i32) -> i64 {
    let n = prices.len();
    let k = k as usize;
    let half = k / 2;
    
    // Base profit (no modification)
    let base: i64 = strategy.iter().zip(prices.iter())
      .map(|(&s, &p)| s as i64 * p as i64).sum();
    
    // With modification at position start (0-indexed):
    // Indices [start, start+half-1] set to 0 (hold)
    // Indices [start+half, start+k-1] set to 1 (sell)
    // Gain from hold part: (0 - strategy[i]) * prices[i] for i in [start, start+half-1]
    // Gain from sell part: (1 - strategy[i]) * prices[i] for i in [start+half, start+k-1]
    //
    // total = base + sum_{i=start}^{start+half-1} (-strategy[i]) * prices[i]
    //               + sum_{i=start+half}^{start+k-1} (1 - strategy[i]) * prices[i]
    //
    // Precompute prefix sums for:
    // A[i] = (-strategy[i]) * prices[i]  (contribution from "zeroing" strategy)
    // B[i] = (1 - strategy[i]) * prices[i]  (contribution from setting to 1)
    
    let a: Vec<i64> = (0..n).map(|i| (-strategy[i] as i64) * prices[i] as i64).collect();
    let b: Vec<i64> = (0..n).map(|i| (1 - strategy[i] as i64) * prices[i] as i64).collect();
    
    let mut prefix_a = vec![0i64; n + 1];
    let mut prefix_b = vec![0i64; n + 1];
    for i in 0..n {
      prefix_a[i+1] = prefix_a[i] + a[i];
      prefix_b[i+1] = prefix_b[i] + b[i];
    }
    
    let range_sum_a = |l: usize, r: usize| prefix_a[r+1] - prefix_a[l]; // inclusive
    let range_sum_b = |l: usize, r: usize| prefix_b[r+1] - prefix_b[l];
    
    let mut best = base;
    for start in 0..=(n - k) {
      let gain = range_sum_a(start, start + half - 1)
           + range_sum_b(start + half, start + k - 1);
      best = best.max(base + gain);
    }
    best
  }
}