Question

2913. 子数组不同元素数目的平方和 I

English Version

题目描述

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

定义 nums 一个子数组的 不同计数 值如下：

• 令 nums[i..j] 表示 nums 中所有下标在 i 到 j 范围内的元素构成的子数组（满足 0 <= i <= j < nums.length ），那么我们称子数组 nums[i..j] 中不同值的数目为 nums[i..j] 的不同计数。

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

由于答案可能会很大，请你将它对 109 + 7 取余 后返回。

子数组指的是一个数组里面一段连续 非空 的元素序列。

 

示例 1：

输入：nums = [1,2,1]
输出：15
解释：六个子数组分别为：
[1]: 1 个互不相同的元素。
[2]: 1 个互不相同的元素。
[1]: 1 个互不相同的元素。
[1,2]: 2 个互不相同的元素。
[2,1]: 2 个互不相同的元素。
[1,2,1]: 2 个互不相同的元素。
所有不同计数的平方和为 12 + 12 + 12 + 22 + 22 + 22 = 15 。

示例 2：

输入：nums = [2,2]
输出：3
解释：三个子数组分别为：
[2]: 1 个互不相同的元素。
[2]: 1 个互不相同的元素。
[2,2]: 1 个互不相同的元素。
所有不同计数的平方和为 12 + 12 + 12 = 3 。

 

提示：

• 1 <= nums.length <= 100
• 1 <= nums[i] <= 100

解法

方法一：枚举

我们可以枚举子数组的左端点下标 $i$，对于每个 $i$，我们在 $[i, n)$ 的范围内枚举子数组的右端点下标 $j$，并统计 $nums[j]$ 的值，将其加入到集合 $s$ 中，记 $s$ 的大小为 $cnt$，那么 $nums[i..j]$ 的不同计数为 $cnt$，将其平方后加入到答案中。

枚举结束后，返回答案即可。

时间复杂度 $O(n^2)$，空间复杂度 $O(n)$。其中 $n$ 是数组 $nums$ 的长度。

class Solution:
  def sumCounts(self, nums: List[int]) -> int:
    ans, n = 0, len(nums)
    for i in range(n):
      s = set()
      for j in range(i, n):
        s.add(nums[j])
        ans += len(s) * len(s)
    return ans

class Solution {
  public int sumCounts(List<Integer> nums) {
    int ans = 0;
    int n = nums.size();
    for (int i = 0; i < n; ++i) {
      int[] s = new int[101];
      int cnt = 0;
      for (int j = i; j < n; ++j) {
        if (++s[nums.get(j)] == 1) {
          ++cnt;
        }
        ans += cnt * cnt;
      }
    }
    return ans;
  }
}

class Solution {
public:
  int sumCounts(vector<int>& nums) {
    int ans = 0;
    int n = nums.size();
    for (int i = 0; i < n; ++i) {
      int s[101]{};
      int cnt = 0;
      for (int j = i; j < n; ++j) {
        if (++s[nums[j]] == 1) {
          ++cnt;
        }
        ans += cnt * cnt;
      }
    }
    return ans;
  }
};

func sumCounts(nums []int) (ans int) {
  for i := range nums {
    s := [101]int{}
    cnt := 0
    for _, x := range nums[i:] {
      s[x]++
      if s[x] == 1 {
        cnt++
      }
      ans += cnt * cnt
    }
  }
  return
}

function sumCounts(nums: number[]): number {
  let ans = 0;
  const n = nums.length;
  for (let i = 0; i < n; ++i) {
    const s: number[] = Array(101).fill(0);
    let cnt = 0;
    for (const x of nums.slice(i)) {
      if (++s[x] === 1) {
        ++cnt;
      }
      ans += cnt * cnt;
    }
  }
  return ans;
}