Skip to content

3473. Sum of K Subarrays With Length at Least M 👍

Approach 1: Top-down

  • Time: $O(nk)$
  • Space: $O(nk)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class Solution {
 public:
  int maxSum(vector<int>& nums, int k, int m) {
    const int n = nums.size();
    vector<int> prefix(n + 1);
    vector<vector<vector<int>>> mem(
        n, vector<vector<int>>(2, vector<int>(k + 1, -1)));
    partial_sum(nums.begin(), nums.end(), prefix.begin() + 1);
    return maxSum(nums, 0, /*ongoing=*/0, k, m, prefix, mem);
  }

 private:
  static constexpr int kInf = 20'000'000;

  int maxSum(const vector<int>& nums, int i, int ongoing, int k, const int& m,
             const vector<int>& prefix, vector<vector<vector<int>>>& mem) {
    if (k < 0)
      return -kInf;
    if (i == nums.size())
      return k == 0 ? 0 : -kInf;
    if (mem[i][ongoing][k] != -1)
      return mem[i][ongoing][k];
    if (ongoing == 1)
      // 1. End the current subarray (transition to state 0, same index i)
      // 2. Extend the current subarray by picking nums[i] and move to i + 1.
      return mem[i][1][k] =
                 max(maxSum(nums, i, 0, k, m, prefix, mem),
                     maxSum(nums, i + 1, 1, k, m, prefix, mem) + nums[i]);
    // ongoing == 0
    // 1. Skip nums[i].
    // 2. Pick nums[i..i + m - 1] (only if k > 0 and there're enough elements).
    int res = maxSum(nums, i + 1, 0, k, m, prefix, mem);
    if (i + m <= nums.size())  // If we have enough elements for a new segment
      res = max(res, maxSum(nums, i + m, 1, k - 1, m, prefix, mem) +
                         (prefix[i + m] - prefix[i]));
    return mem[i][0][k] = res;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Solution {
  public int maxSum(int[] nums, int k, int m) {
    final int n = nums.length;
    int[] prefix = new int[n + 1];
    Integer[][][] mem = new Integer[n][2][k + 1];
    for (int i = 0; i < n; i++)
      prefix[i + 1] = prefix[i] + nums[i];
    return maxSum(nums, 0, 0, k, m, prefix, mem);
  }

  private static final int INF = 20_000_000;

  private int maxSum(int[] nums, int i, int ongoing, int k, int m, int[] prefix,
                     Integer[][][] mem) {
    if (k < 0)
      return -INF;
    if (i == nums.length)
      return k == 0 ? 0 : -INF;
    if (mem[i][ongoing][k] != null)
      return mem[i][ongoing][k];
    if (ongoing == 1)
      // 1. End the current subarray (transition to state 0, same index i)
      // 2. Extend the current subarray by picking nums[i] and move to i + 1.
      return mem[i][1][k] = Math.max(maxSum(nums, i, 0, k, m, prefix, mem),
                                     maxSum(nums, i + 1, 1, k, m, prefix, mem) + nums[i]);
    // ongoing == 0
    // 1. Skip nums[i].
    // 2. Pick nums[i..i + m - 1] (only if k > 0 and there're enough elements).
    int res = maxSum(nums, i + 1, 0, k, m, prefix, mem);
    if (i + m <= nums.length) // If we have enough elements for a new segment
      res = Math.max(res,
                     maxSum(nums, i + m, 1, k - 1, m, prefix, mem) + (prefix[i + m] - prefix[i]));

    return mem[i][0][k] = res;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Solution:
  def maxSum(self, nums: list[int], k: int, m: int) -> int:
    INF = 20_000_000
    n = len(nums)
    prefix = list(itertools.accumulate(nums, initial=0))

    @functools.lru_cache(None)
    def dp(i: int, ongoing: int, k: int) -> int:
      if k < 0:
        return -INF
      if i == n:
        return 0 if k == 0 else -INF
      if ongoing == 1:
        # 1. End the current subarray (transition to state 0, same index i)
        # 2. Extend the current subarray by picking nums[i] and move to i + 1
        return max(dp(i, 0, k),
                   dp(i + 1, 1, k) + nums[i])
      # ongoing == 0
      # 1. Skip nums[i]
      # 2. Pick nums[i:i+m] (only if k > 0 and there're enough elements)
      res = dp(i + 1, 0, k)
      if i + m <= n:  # If we have enough elements for a new segment
        res = max(res,
                  dp(i + m, 1, k - 1) + (prefix[i + m] - prefix[i]))
      return res

    return dp(0, 0, k)

Approach 2: Bottom-up

  • Time: $O(nk)$
  • Space: $O(nk)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
class Solution {
 public:
  int maxSum(vector<int>& nums, int k, int m) {
    constexpr int kInf = 20000000;
    const int n = nums.size();
    vector<int> prefix(n + 1);
    // dp[i][ongoing][r] := the maximum sum of nums[i..n - 1], with `ongoing`
    // indicating if a subarray is currently being extended (1) or not (0), and
    // `r` segments left to choose
    vector<vector<vector<int>>> dp(
        n + 1, vector<vector<int>>(2, vector<int>(k + 1, -kInf)));

    partial_sum(nums.begin(), nums.end(), prefix.begin() + 1);

    // Base case: At the end of the array, if no segments are left, score is 0.
    dp[n][0][0] = dp[n][1][0] = 0;

    for (int i = n - 1; i >= 0; --i)
      for (int rem = 0; rem <= k; ++rem) {
        // When no subarray is ongoing:
        // 1. Skip nums[i].
        dp[i][0][rem] = dp[i + 1][0][rem];
        // 2. Start a new segment of length m (only if rem > 0 and there're
        // enough elements)
        if (rem > 0 && i + m <= n)
          dp[i][0][rem] = max(dp[i][0][rem], dp[i + m][1][rem - 1] +
                                                 (prefix[i + m] - prefix[i]));
        // When a subarray is ongoing:
        // 1. End the current subarray (transition to state 0, same index i)
        // 2. Extend the current subarray by picking nums[i] and move to i + 1
        dp[i][1][rem] = max(dp[i][0][rem], dp[i + 1][1][rem] + nums[i]);
      }

    return dp[0][0][k];
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Solution {
  public int maxSum(int[] nums, int k, int m) {
    final int INF = 20_000_000;
    final int n = nums.length;
    int[] prefix = new int[n + 1];
    // dp[i][ongoing][r] := the maximum sum of nums[i..n - 1], with `ongoing`
    // indicating if a subarray is currently being extended (1) or not (0), and
    // `r` segments left to choose
    int[][][] dp = new int[n + 1][2][k + 1];

    for (int i = 0; i < n; ++i)
      prefix[i + 1] = prefix[i] + nums[i];

    Arrays.stream(dp).forEach(row -> Arrays.stream(row).forEach(col -> Arrays.fill(col, -INF)));

    // Base case: At the end of the array, if no segments are left, score is 0.
    dp[n][0][0] = dp[n][1][0] = 0;

    for (int i = n - 1; i >= 0; --i)
      for (int rem = 0; rem <= k; ++rem) {
        // When no subarray is ongoing:
        // 1. Skip nums[i].
        dp[i][0][rem] = dp[i + 1][0][rem];
        // 2. Start a new segment of length m (only if rem > 0 and there're
        // enough elements)
        if (rem > 0 && i + m <= n)
          dp[i][0][rem] =
              Math.max(dp[i][0][rem], dp[i + m][1][rem - 1] + (prefix[i + m] - prefix[i]));
        // When a subarray is ongoing:
        // 1. End the current subarray (transition to state 0, same index i)
        // 2. Extend the current subarray by picking nums[i] and move to i + 1
        dp[i][1][rem] = Math.max(dp[i][0][rem], dp[i + 1][1][rem] + nums[i]);
      }

    return dp[0][0][k];
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class Solution:
  def maxSum(self, nums: list[int], k: int, m: int) -> int:
    INF = 20_000_000
    n = len(nums)
    prefix = list(itertools.accumulate(nums, initial=0))

    # dp[i][ongoing][r] := the maximum sum of nums[i:], with `ongoing`
    # indicating if a subarray is currently being extended (1) or not (0),
    # and `r` segments left to choose
    dp = [[[-INF] * (k + 1) for _ in range(2)] for _ in range(n + 1)]

    # Base case: At the end of the array, if no segments are left, score is 0
    dp[n][0][0] = dp[n][1][0] = 0

    for i in range(n - 1, -1, -1):
      for rem in range(k + 1):
        # When no subarray is ongoing:
        # 1. Skip nums[i]
        dp[i][0][rem] = dp[i + 1][0][rem]
        # 2. Start a new segment of length m (only if rem > 0 and there're enough elements)
        if rem > 0 and i + m <= n:
          dp[i][0][rem] = max(
              dp[i][0][rem],
              dp[i + m][1][rem - 1] + (prefix[i + m] - prefix[i]))
        # When a subarray is ongoing:
        # 1. End the current subarray (transition to state 0, same index i)
        # 2. Extend the current subarray by picking nums[i] and move to i + 1
        dp[i][1][rem] = max(dp[i][0][rem], dp[i + 1][1][rem] + nums[i])

    return dp[0][0][k]