题目

给你一个下标从 0 开始且长度为 n 的整数数组 nums 。分割 数组 nums 的方案数定义为符合以下两个条件的 pivot 数目：

1 <= pivot < n

nums[0] + nums[1] + ... + nums[pivot - 1] == nums[pivot] + nums[pivot + 1] + ... + nums[n - 1]

同时给你一个整数 k 。你可以将 nums 中 一个 元素变为 k 或 不改变 数组。

请你返回在 至多 改变一个元素的前提下，最多 有多少种方法 分割 nums 使得上述两个条件都满足。

示例 1：输入：nums = [2,-1,2], k = 3 输出：1

解释：一个最优的方案是将 nums[0] 改为 k 。数组变为 [3,-1,2] 。

有一种方法分割数组：

- pivot = 2 ，我们有分割 [3,-1 | 2]：3 + -1 == 2 。

示例 2：输入：nums = [0,0,0], k = 1 输出：2

解释：一个最优的方案是不改动数组。

有两种方法分割数组：

- pivot = 1 ，我们有分割 [0 | 0,0]：0 == 0 + 0 。

- pivot = 2 ，我们有分割 [0,0 | 0]: 0 + 0 == 0 。

示例 3：输入：nums = [22,4,-25,-20,-15,15,-16,7,19,-10,0,-13,-14], k = -33 输出：4

解释：一个最优的方案是将 nums[2] 改为 k 。数组变为 [22,4,-33,-20,-15,15,-16,7,19,-10,0,-13,-14] 。

有四种方法分割数组。

提示：n == nums.length

2 <= n <= 105

-105 <= k, nums[i] <= 105

解题思路分析

1、前缀和+哈希；时间复杂度O(n)，空间复杂度O(n)

func waysToPartition(nums []int, k int) int {
   res := 0
   n := len(nums)
   sum := 0
   prev := make(map[int]int) // 前缀和
   arr := make([]int, n+1)
   for i := 0; i < n; i++ {
      arr[i+1] = arr[i] + nums[i]
      sum = sum + nums[i]
      if i != n-1 {
         prev[sum]++
      }
   }
   if sum%2 == 0 {
      res = prev[sum/2] // 不改变的结果
   }
   suf := make(map[int]int) // 后缀和
   sufSum := 0
   temp := sum
   for i := n - 1; i >= 0; i-- { // 枚举每一位修改后的结果
      target := sum - nums[i] + k // 替换后的总和
      temp = temp - nums[i]
      prev[temp]-- // 前缀和 减一
      sufSum = sufSum + k
      suf[sufSum]++
      if target%2 == 0 {
         res = max(res, prev[target/2]+suf[target/2])
      }
      suf[sufSum]--
      sufSum = sufSum - k + nums[i]
      suf[sufSum]++
   }
   return res
}

func max(a, b int) int {
   if a > b {
      return a
   }
   return b
}

2、前缀和；时间复杂度O(n)，空间复杂度O(n)

func waysToPartition(nums []int, k int) int {
   res := 0
   n := len(nums)
   prev := make(map[int]int) // 前缀和
   arr := make([]int, n)
   arr[0] = nums[0]
   for i := 1; i < n; i++ {
      arr[i] = arr[i-1] + nums[i]
      prev[arr[i-1]]++
   }
   sum := arr[n-1]
   if sum%2 == 0 {
      res = prev[sum/2] // 不改变的结果
   }
   m := make(map[int]int)
   for i := 0; i < len(arr); i++ {
      diff := k - nums[i]
      if (sum+diff)%2 == 0 {
         res = max(res, prev[(sum-diff)/2]+m[(sum+diff)/2])
      }
      m[arr[i]]++    // 左侧+1
      prev[arr[i]]-- // 右侧-1
   }
   return res
}

func max(a, b int) int {
   if a > b {
      return a
   }
   return b
}

总结

Hard题目，前缀和题目