Problem

You are given a 0-indexed integer array nums.

The distinct count of a subarray of nums is defined as:

  • Let nums[i..j] be a subarray of nums consisting of all the indices from i to j such that 0 <= i <= j < nums.length. Then the number of distinct values in nums[i..j] is called the distinct count of nums[i..j].

Return _the sum of thesquares of distinct counts of all subarrays of _nums.

A subarray is a contiguous non-empty sequence of elements within an array.

Examples

Example 1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

Input: nums = [1,2,1]
Output: 15
Explanation: Six possible subarrays are:
[1]: 1 distinct value
[2]: 1 distinct value
[1]: 1 distinct value
[1,2]: 2 distinct values
[2,1]: 2 distinct values
[1,2,1]: 2 distinct values
The sum of the squares of the distinct counts in all subarrays is equal to 12 + 12 + 12 + 22 + 22 + 22 = 15.

Example 2

1
2
3
4
5
6
7

Input: nums = [1,1]
Output: 3
Explanation: Three possible subarrays are:
[1]: 1 distinct value
[1]: 1 distinct value
[1,1]: 1 distinct value
The sum of the squares of the distinct counts in all subarrays is equal to 12 + 12 + 12 = 3.

Constraints

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

Solution

Method 1 - Brute Force with Hash Set

Intuition

Since the array is small (n ≤ 100), we can check all subarrays and count the number of distinct elements using a set, then sum the squares.

Approach

  1. For each subarray, use a set to count the number of distinct elements.
  2. For each subarray, add the square of the number of distinct elements to the answer.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

#include <vector>
#include <unordered_set>
using namespace std;
class Solution {
public:
    int sumCounts(vector<int>& nums) {
        int n = nums.size(), ans = 0;
        for (int l = 0; l < n; ++l) {
            unordered_set<int> s;
            for (int r = l; r < n; ++r) {
                s.insert(nums[r]);
                ans += s.size() * s.size();
            }
        }
        return ans;
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

func sumCounts(nums []int) int {
    n := len(nums)
    ans := 0
    for l := 0; l < n; l++ {
        s := map[int]bool{}
        for r := l; r < n; r++ {
            s[nums[r]] = true
            ans += len(s) * len(s)
        }
    }
    return ans
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

import java.util.*;
class Solution {
    public int sumCounts(int[] nums) {
        int n = nums.length, ans = 0;
        for (int l = 0; l < n; ++l) {
            Set<Integer> s = new HashSet<>();
            for (int r = l; r < n; ++r) {
                s.add(nums[r]);
                ans += s.size() * s.size();
            }
        }
        return ans;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

fun sumCounts(nums: IntArray): Int {
    val n = nums.size
    var ans = 0
    for (l in 0 until n) {
        val s = mutableSetOf<Int>()
        for (r in l until n) {
            s.add(nums[r])
            ans += s.size * s.size
        }
    }
    return ans
}
1
2
3
4
5
6
7
8
9

def sumCounts(nums):
    n = len(nums)
    ans = 0
    for l in range(n):
        s = set()
        for r in range(l, n):
            s.add(nums[r])
            ans += len(s) * len(s)
    return ans
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

use std::collections::HashSet;
pub fn sum_counts(nums: Vec<i32>) -> i32 {
    let n = nums.len();
    let mut ans = 0;
    for l in 0..n {
        let mut s = HashSet::new();
        for r in l..n {
            s.insert(nums[r]);
            ans += (s.len() * s.len()) as i32;
        }
    }
    ans
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

function sumCounts(nums: number[]): number {
    const n = nums.length;
    let ans = 0;
    for (let l = 0; l < n; l++) {
        const s = new Set<number>();
        for (let r = l; r < n; r++) {
            s.add(nums[r]);
            ans += s.size * s.size;
        }
    }
    return ans;
}

Complexity

  • ⏰ Time complexity: O(n^3)
  • 🧺 Space complexity: O(n)