Hard

题目描述

给你一个下标从 0 开始的整数数组 nums。

nums 中一个子数组的 不同计数 定义为：

设 nums[i..j] 表示 nums 中所有下标从 i 到 j 的元素构成的子数组，其中 0 <= i <= j < nums.length，那么 nums[i..j] 中不同值的数目称为 nums[i..j] 的不同计数。

返回 nums 所有子数组的不同计数的 平方和。

由于答案可能很大，将它对 10^9 + 7 取余 后返回。

子数组是数组中一个连续 非空 的元素序列。

示例 1：

输入：nums = [1,2,1]
输出：15
解释：六个可能的子数组分别为：
[1]: 1 个不同值
[2]: 1 个不同值  
[1]: 1 个不同值
[1,2]: 2 个不同值
[2,1]: 2 个不同值
[1,2,1]: 2 个不同值
所有子数组不同计数的平方和等于 1² + 1² + 1² + 2² + 2² + 2² = 15

示例 2：

输入：nums = [2,2]
输出：3
解释：三个可能的子数组分别为：
[2]: 1 个不同值
[2]: 1 个不同值
[2,2]: 1 个不同值
所有子数组不同计数的平方和等于 1² + 1² + 1² = 3

提示：

  • 1 <= nums.length <= 10^5
  • 1 <= nums[i] <= 10^5

解题思路

这道题要计算所有子数组不同元素数目的平方和。直接枚举所有子数组会导致 O(n³) 的时间复杂度，无法通过。

核心思路： 我们考虑动态规划的方法。设 dp[i] 表示以位置 i 结尾的所有子数组不同元素数目的平方和。

关键观察：当我们从位置 i-1 扩展到位置 i 时，新增了元素 nums[i]。对于每个以 i-1 结尾的子数组，扩展到位置 i 后：

  • 如果 nums[i] 之前没有出现过，不同元素数目加1
  • 如果 nums[i] 在位置 j 出现过，那么从位置 j+1 到 i-1 的子数组扩展后不同元素数目加1，而从更早位置开始的子数组不同元素数目不变

状态转移： 设 sum[i] 表示以位置 i 结尾的所有子数组的不同元素数目之和。

  • 如果 nums[i] 之前没出现过：sum[i] = sum[i-1] + (i+1)，dp[i] = dp[i-1] + 2*sum[i-1] + (i+1)
  • 如果 nums[i] 在位置 last 出现过：需要减去重复计算的部分

为了高效处理区间和与区间平方和的更新，我们使用线段树或树状数组来维护这些信息。

推荐解法： 使用动态规划 + 线段树，时间复杂度 O(n log n)。

代码实现

class Solution {
public:
    const int MOD = 1e9 + 7;
    
    int sumCounts(vector<int>& nums) {
        int n = nums.size();
        vector<int> last(100001, -1);
        
        long long ans = 0;
        long long sum = 0, sum2 = 0;
        
        for (int i = 0; i < n; i++) {
            int prev = last[nums[i]];
            last[nums[i]] = i;
            
            long long newSum = (sum + i + 1 - prev) % MOD;
            long long newSum2 = (sum2 + 2 * (sum * (i - prev) % MOD) + 
                                (long long)(i - prev) * (i - prev) % MOD) % MOD;
            
            sum = newSum;
            sum2 = newSum2;
            ans = (ans + sum2) % MOD;
        }
        
        return ans;
    }
};

class Solution:
    def sumCounts(self, nums: List[int]) -> int:
        MOD = 10**9 + 7
        n = len(nums)
        last = {}
        
        ans = 0
        sum_val = 0
        sum2_val = 0
        
        for i in range(n):
            prev = last.get(nums[i], -1)
            last[nums[i]] = i
            
            diff = i - prev
            new_sum = (sum_val + diff) % MOD
            new_sum2 = (sum2_val + 2 * sum_val * diff + diff * diff) % MOD
            
            sum_val = new_sum
            sum2_val = new_sum2
            ans = (ans + sum2_val) % MOD
            
        return ans

public class Solution {
    public int SumCounts(int[] nums) {
        const int MOD = 1000000007;
        int n = nums.Length;
        Dictionary<int, int> last = new Dictionary<int, int>();
        
        long ans = 0;
        long sum = 0, sum2 = 0;
        
        for (int i = 0; i < n; i++) {
            int prev = last.GetValueOrDefault(nums[i], -1);
            last[nums[i]] = i;
            
            long diff = i - prev;
            long newSum = (sum + diff) % MOD;
            long newSum2 = (sum2 + 2 * sum * diff % MOD + diff * diff % MOD) % MOD;
            
            sum = newSum;
            sum2 = newSum2;
            ans = (ans + sum2) % MOD;
        }
        
        return (int)ans;
    }
}

var sumCounts = function(nums) {
    const MOD = 1e9 + 7;
    const n = nums.length;
    const last = new Map();
    
    let ans = 0;
    let sum = 0, sum2 = 0;
    
    for (let i = 0; i < n; i++) {
        const prev = last.get(nums[i]) ?? -1;
        last.set(nums[i], i);
        
        const diff = i - prev;
        const newSum = (sum + diff) % MOD;
        const newSum2 = (sum2 + 2 * sum * diff % MOD + diff * diff % MOD) % MOD;
        
        sum = newSum;
        sum2 = newSum2;
        ans = (ans + sum2) % MOD;
    }
    
    return ans;
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

复杂度大小
时间复杂度O(n)
空间复杂度O(k)，其中 k 是不同元素的数目