Hard

题目描述

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

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

同时给你一个整数 k。你可以选择将 nums 中一个元素的值更改为 k，或者不更改数组。

请你返回在更改至多一个元素的前提下，最多能够让多少个 pivot 下标满足上述两个条件。

示例 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 <= 10^5
  • -10^5 <= k, nums[i] <= 10^5

解题思路

这道题需要分析在最多改变一个元素的情况下，如何最大化有效分割点的数量。

思路分析：

首先理解分割条件：对于分割点 pivot，左侧元素和等于右侧元素和。设 left[i] 为前 i 个元素和，right[i] 为从第 i 个元素开始的后缀和，则条件为 left[pivot-1] == right[pivot]。

我们可以用前缀和来表示：prefix[pivot-1] == total - prefix[pivot-1]，即 2 * prefix[pivot-1] == total。

核心观察： 当我们将 nums[i] 改为 k 时，差值 delta = k - nums[i] 会影响所有包含该元素的前缀和。

算法步骤：

  1. 计算原始数组的所有前缀和，找出原始的有效分割点数量
  2. 对于每个可能的修改位置 i，计算修改后能增加的有效分割点数量
  3. 当修改 nums[i] 时：
    • 对于 pivot <= i 的分割点，左侧和不受影响
    • 对于 pivot > i 的分割点，左侧和增加 delta
  4. 使用哈希表记录每个差值出现的次数，快速计算修改后的收益

具体来说，我们维护两个哈希表：

  • left_count：记录在当前位置左侧的差值分布
  • right_count：记录在当前位置右侧的差值分布

遍历每个位置时，计算如果在该位置修改，能让多少个分割点变为有效。

代码实现

class Solution {
public:
    int waysToPartition(vector<int>& nums, int k) {
        int n = nums.size();
        vector<long long> prefix(n + 1, 0);
        
        for (int i = 0; i < n; i++) {
            prefix[i + 1] = prefix[i] + nums[i];
        }
        
        long long total = prefix[n];
        
        // 计算原始数组的有效分割数
        unordered_map<long long, int> right_count;
        for (int pivot = 1; pivot < n; pivot++) {
            long long diff = 2 * prefix[pivot] - total;
            right_count[diff]++;
        }
        
        int result = right_count[0];  // 不修改的情况
        
        unordered_map<long long, int> left_count;
        
        // 遍历每个可能的修改位置
        for (int i = 0; i < n; i++) {
            long long delta = k - nums[i];
            
            // 修改nums[i]后，能让多少个分割点变有效
            int ways = left_count[-delta] + right_count[delta];
            result = max(result, ways);
            
            // 更新left_count和right_count
            if (i + 1 < n) {
                long long diff = 2 * prefix[i + 1] - total;
                right_count[diff]--;
                if (right_count[diff] == 0) {
                    right_count.erase(diff);
                }
                left_count[diff]++;
            }
        }
        
        return result;
    }
};

class Solution:
    def waysToPartition(self, nums: List[int], k: int) -> int:
        n = len(nums)
        prefix = [0] * (n + 1)
        
        for i in range(n):
            prefix[i + 1] = prefix[i] + nums[i]
        
        total = prefix[n]
        
        # 计算原始数组的有效分割数
        right_count = defaultdict(int)
        for pivot in range(1, n):
            diff = 2 * prefix[pivot] - total
            right_count[diff] += 1
        
        result = right_count[0]  # 不修改的情况
        
        left_count = defaultdict(int)
        
        # 遍历每个可能的修改位置
        for i in range(n):
            delta = k - nums[i]
            
            # 修改nums[i]后，能让多少个分割点变有效
            ways = left_count[-delta] + right_count[delta]
            result = max(result, ways)
            
            # 更新left_count和right_count
            if i + 1 < n:
                diff = 2 * prefix[i + 1] - total
                right_count[diff] -= 1
                if right_count[diff] == 0:
                    del right_count[diff]
                left_count[diff] += 1
        
        return result

public class Solution {
    public int WaysToPartition(int[] nums, int k) {
        int n = nums.Length;
        long[] prefix = new long[n + 1];
        
        for (int i = 0; i < n; i++) {
            prefix[i + 1] = prefix[i] + nums[i];
        }
        
        long total = prefix[n];
        
        // 计算原始数组的有效分割数
        var rightCount = new Dictionary<long, int>();
        for (int pivot = 1; pivot < n; pivot++) {
            long diff = 2 * prefix[pivot] - total;
            rightCount[diff] = rightCount.GetValueOrDefault(diff, 0) + 1;
        }
        
        int result = rightCount.GetValueOrDefault(0, 0);  // 不修改的情况
        
        var leftCount = new Dictionary<long, int>();
        
        // 遍历每个可能的修改位置
        for (int i = 0; i < n; i++) {
            long delta = k - nums[i];
            
            // 修改nums[i]后，能让多少个分割点变有效
            int ways = leftCount.GetValueOrDefault(-delta, 0) + rightCount.GetValueOrDefault(delta, 0);
            result = Math.Max(result, ways);
            
            // 更新leftCount和rightCount
            if (i + 1 < n) {
                long diff = 2 * prefix[i + 1] - total;
                rightCount[diff] = rightCount.GetValueOrDefault(diff, 0) - 1;
                if (rightCount[diff] == 0) {
                    rightCount.Remove(diff);
                }
                leftCount[diff] = leftCount.GetValueOrDefault(diff, 0) + 1;
            }
        }
        
        return result;
    }
}

var waysToPartition = function(nums, k) {
    const n = nums.length;
    const prefixSum = new Array(n + 1).fill(0);
    
    for (let i = 0; i < n; i++) {
        prefixSum[i + 1] = prefixSum[i] + nums[i];
    }
    
    const total = prefixSum[n];
    let originalWays = 0;
    
    for (let pivot = 1; pivot < n; pivot++) {
        if (prefixSum[pivot] === total - prefixSum[pivot]) {
            originalWays++;
        }
    }
    
    let maxWays = originalWays;
    
    for (let changeIdx = 0; changeIdx < n; changeIdx++) {
        const diff = k - nums[changeIdx];
        const newTotal = total + diff;
        let ways = 0;
        
        for (let pivot = 1; pivot < n; pivot++) {
            let leftSum = prefixSum[pivot];
            let rightSum = total - prefixSum[pivot];
            
            if (changeIdx < pivot) {
                leftSum += diff;
            } else {
                rightSum += diff;
            }
            
            if (leftSum === rightSum) {
                ways++;
            }
        }
        
        maxWays = Math.max(maxWays, ways);
    }
    
    return maxWays;
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型值
时间复杂度O(n)
空间复杂度O(n)

其中 n 是数组长度。我们需要遍历数组一次计算前缀和，再遍历一次处理每个修改位置，哈希表的空间复杂度最多为 O(n)。

相关题目