Subarrays Distinct Element Sum of Squares I
EasyPrompt
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 ofnumsconsisting of all the indices fromitojsuch that0 <= i <= j < nums.length. Then the number of distinct values innums[i..j]is called the distinct count ofnums[i..j].
Return the sum of the squares of distinct counts of all subarrays of nums.
A subarray is a contiguous non-empty sequence of elements within an array.
Example 1:
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:
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 <= 1001 <= nums[i] <= 100
Approaches
2 approaches with complexity analysis and trade-offs.
This approach is the most straightforward and directly follows the problem description. It involves generating every possible subarray of the input array nums. For each of these subarrays, we iterate through its elements to count the number of distinct values using a HashSet. Finally, we square this count and add it to a running total.
Algorithm
- Initialize a variable
totalSumto 0. - Use a pair of nested loops with indices
iandjto generate all possible subarrays. The outer loopiruns from0ton-1(start index), and the inner loopjruns fromiton-1(end index). - For each subarray
nums[i..j], create a newHashSetto store its unique elements. - Use a third nested loop with index
kfromitojto iterate through the current subarray. - Add each element
nums[k]to theHashSet. - After iterating through the subarray, the size of the
HashSetgives the count of distinct elements. - Square this count and add it to
totalSum. - After all subarrays have been processed, return
totalSum.
Walkthrough
The algorithm employs three nested loops. The first two loops, with indices i and j, are used to define the start and end of every subarray nums[i..j]. For each such subarray, a third loop with index k is used to traverse its elements. Inside this innermost loop, we use a HashSet to keep track of the unique elements encountered. The size of the set after the loop gives the distinct count. This count is then squared and accumulated into a final sum. This process is repeated for all n * (n + 1) / 2 subarrays.
import java.util.HashSet;import java.util.Set; class Solution { public int sumCounts(int[] nums) { int n = nums.length; int totalSum = 0; for (int i = 0; i < n; i++) { for (int j = i; j < n; j++) { // Subarray is nums[i..j] Set<Integer> distinctElements = new HashSet<>(); for (int k = i; k <= j; k++) { distinctElements.add(nums[k]); } int distinctCount = distinctElements.size(); totalSum += distinctCount * distinctCount; } } return totalSum; }}Complexity
Time
O(N^3), where N is the length of `nums`. There are three nested loops. The outer two loops select a subarray (O(N^2) subarrays), and the inner loop iterates through the subarray's elements (up to O(N) elements), leading to a cubic time complexity.
Space
O(N), where N is the length of `nums`. The `HashSet` can store up to N distinct elements in the worst case for a single subarray.
Trade-offs
Pros
Simple to understand and implement.
Directly translates the problem statement into code without complex logic.
Cons
Highly inefficient due to its cubic time complexity.
Performs a lot of redundant work. For each subarray, it recalculates the distinct elements from scratch, even for overlapping subarrays.
This approach improves upon the brute-force method by avoiding redundant calculations. Instead of re-calculating the distinct count for each subarray from scratch, we can efficiently update the count as we extend the subarray one element at a time. This reduces the overall time complexity from cubic to quadratic.
Algorithm
- Initialize a variable
totalSumto 0. - Use an outer loop with index
ifrom0ton-1to fix the starting point of subarrays. - For each starting index
i, create a newHashSet. - Use an inner loop with index
jfromiton-1. This loop extends the subarray to the right. - In each iteration of the inner loop, add the element
nums[j]to theHashSet. The set now holds the distinct elements for the subarraynums[i..j]. - The size of the
HashSetis the distinct count for the current subarray. - Square this count and add it to
totalSum. - After the loops complete, return
totalSum.
Walkthrough
We use two nested loops. The outer loop (with index i) fixes the starting point of our subarrays. For each starting point i, we initialize a HashSet. The inner loop (with index j) starts from i and iterates to the end of the array. In each step of the inner loop, we consider the subarray nums[i..j]. We simply add the new element nums[j] to the HashSet we created for the starting index i. The set now correctly represents the distinct elements for nums[i..j]. The size of the set gives us the distinct count, which we square and add to our total. By reusing the HashSet for all subarrays starting at i, we reduce the complexity of finding the distinct count for each subarray to O(1) on average.
import java.util.HashSet;import java.util.Set; class Solution { public int sumCounts(int[] nums) { int n = nums.length; int totalSum = 0; for (int i = 0; i < n; i++) { Set<Integer> distinctElements = new HashSet<>(); for (int j = i; j < n; j++) { // Extend the subarray nums[i..j-1] to nums[i..j] distinctElements.add(nums[j]); int distinctCount = distinctElements.size(); totalSum += distinctCount * distinctCount; } } return totalSum; }}Complexity
Time
O(N^2), where N is the length of `nums`. We have two nested loops. The inner operation (adding to a set and getting its size) takes O(1) on average. This is the most common and accepted solution for this type of problem with small constraints.
Space
O(N), where N is the length of `nums`. The `HashSet` is re-initialized for each starting index `i` and can store up to N elements in the worst case.
Trade-offs
Pros
Significantly more efficient than the O(N^3) approach.
Simple to implement and fast enough for the given constraints.
Eliminates redundant work by incrementally building the set of distinct elements.
Cons
While efficient for the given constraints, it would be too slow for much larger inputs (e.g., N > 10^4).
Solutions
Solution
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; }};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 ansVideo walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.