Skip to content

3077. Maximum Strength of K Disjoint Subarrays

  • 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:
  long long maximumStrength(vector<int>& nums, int k) {
    vector<vector<vector<long>>> mem(
        nums.size(), vector<vector<long>>(k + 1, vector<long>(2, -1)));
    return maximumStrength(nums, 0, k, /*fresh=*/true, mem);
  }

 private:
  static constexpr long kMin = LONG_MIN / 2;

  // Returns the maximum strength of nums[i..n) with k operations left, where
  // `fresh` means we're starting a new subarray.
  long maximumStrength(const vector<int>& nums, int i, int k, bool fresh,
                       vector<vector<vector<long>>>& mem) {
    if (nums.size() - i < k)
      return kMin;
    if (k == 0)
      return 0;
    if (i == nums.size())
      return k == 0 ? 0 : kMin;
    if (mem[i][k][fresh] != -1)
      return mem[i][k][fresh];
    // If it's not fresh, we can't skip the current number and consider it as a
    // fresh start, since the case where it's fresh is already covered by
    // `includeAndFreshStart`.
    const long skip = fresh ? maximumStrength(nums, i + 1, k, true, mem) : kMin;
    const long gain = (k % 2 == 0 ? -1 : 1) * static_cast<long>(nums[i]) * k;
    const long includeAndContinue =
        maximumStrength(nums, i + 1, k, false, mem) + gain;
    const long includeAndFreshStart =
        maximumStrength(nums, i + 1, k - 1, true, mem) + gain;
    return mem[i][k][fresh] =
               max(skip, max(includeAndContinue, includeAndFreshStart));
  }
};

 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 {
  public long maximumStrength(int[] nums, int k) {
    Long[][][] mem = new Long[nums.length][k + 1][2];
    return maximumStrength(nums, 0, k, /*fresh=*/true, mem);
  }

  private static final long MIN = Long.MIN_VALUE / 2;

  // Returns the maximum strength of nums[i..n) with k operations left, where
  // `fresh` means we're starting a new subarray.
  private long maximumStrength(int[] nums, int i, int k, boolean fresh, Long[][][] mem) {
    if (nums.length - i < k)
      return MIN;
    if (k == 0)
      return 0;
    if (i == nums.length)
      return k == 0 ? 0 : MIN;
    if (mem[i][k][fresh ? 1 : 0] != null)
      return mem[i][k][fresh ? 1 : 0];
    // If it's not fresh, we can't skip the current number and consider it as a
    // fresh start, since the case where it's fresh is already covered by
    // `includeAndFreshStart`.
    final long skip = fresh ? maximumStrength(nums, i + 1, k, true, mem) : MIN;
    final long gain = (k % 2 == 0 ? -1 : 1) * 1L * nums[i] * k;
    final long includeAndContinue = maximumStrength(nums, i + 1, k, false, mem) + gain;
    final long includeAndFreshStart = maximumStrength(nums, i + 1, k - 1, true, mem) + gain;
    return mem[i][k][fresh ? 1 : 0] =
               Math.max(skip, Math.max(includeAndContinue, includeAndFreshStart));
  }
}

 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
class Solution:
  def maximumStrength(self, nums: list[int], k: int) -> int:

    @functools.lru_cache(None)
    def dp(i: int, k: int, fresh: bool) -> int:
      """
      Returns the maximum strength of nums[i..n) with k operations left, where
      `fresh` means we're starting a new subarray.
      """
      if len(nums) - i < k:
        return -math.inf
      if k == 0:
        return 0
      if i == len(nums):
        return 0 if k == 0 else -math.inf
      # If it's not fresh, we can't skip the current number and consider it as a
      # fresh start, since the case where it's fresh is already covered by
      # `includeAndFreshStart`.
      skip = dp(i + 1, k, True) if fresh else -math.inf
      gain = (-1 if k % 2 == 0 else 1) * nums[i] * k
      includeAndContinue = dp(i + 1, k, False) + gain
      includeAndFreshStart = dp(i + 1, k - 1, True) + gain
      return max(skip, includeAndContinue, includeAndFreshStart)

    return dp(0, k, True)