Skip to main content
Talentd

Subarrays Distinct Element Sum of Squares II - Solution & Explanation

HardArrayDynamic ProgrammingBinary Indexed TreeSegment Tree6 min readAsked at: Amazon
Practice this problem

Problem Statement

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 the squares of distinct counts of all subarrays of nums.

Since the answer may be very large, return it modulo 109 + 7.

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 = [2,2]
Output: 3
Explanation: Three possible subarrays are:
[2]: 1 distinct value
[2]: 1 distinct value
[2,2]: 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 <= 105
  • 1 <= nums[i] <= 105

Approach Overview

Problem Overview: Given an array nums, evaluate every subarray and count how many distinct elements it contains. Square that count and sum the result across all subarrays. The challenge is avoiding the naive O(n²) enumeration of distinct counts.

Approach 1: Brute Force with Set Tracking (O(n²) time, O(n) space)

Start a subarray at index i and expand the right boundary one step at a time. Maintain a set or frequency map to track distinct elements as you iterate j = i..n-1. After each extension, compute the current number of unique elements and add distinct² to the answer. This approach directly simulates the definition of the problem and is easy to reason about, but repeated scans of overlapping subarrays make it quadratic.

Approach 2: Sliding Window + Contribution Tracking (O(n log n) time, O(n) space)

The optimized solution tracks how each new element affects the number of distinct elements across multiple subarrays simultaneously. Maintain the last occurrence of each value and compute how many subarrays ending at the current index gain a new distinct element. Instead of recomputing counts for every subarray, accumulate contributions using a data structure such as a Fenwick Tree or Segment Tree. These structures efficiently update ranges and query aggregated values while processing the array left to right. Each element updates the contribution range where it becomes a new distinct value, and the squared distinct counts are aggregated incrementally.

This approach relies on range updates and prefix queries, which makes data structures from Binary Indexed Tree and Segment Tree especially useful. The state transitions resemble a dynamic contribution model often seen in advanced Dynamic Programming problems on subarrays.

Recommended for interviews: Start by explaining the brute force approach to show you understand the definition of the problem. Then move to the optimized contribution-based solution using Fenwick or segment trees. Interviewers typically expect the O(n log n) approach because it demonstrates strong control over range updates, prefix queries, and subarray contribution analysis.

Approach 1: Brute Force Approach

This method involves generating all possible subarrays of the given array and calculating the number of distinct elements in each subarray. The distinct count for each subarray is then squared and summed up to get the final result.

The Python solution initializes a total_sum to zero. It then iterates over all possible start points of subarrays. For each start, it calculates the distinct count while moving to different end points, storing encountered numbers in a set. The square of this distinct count is added to the total_sum, which is returned after taking modulo 109 + 7.

Code

PythonC++

Python

C++

Complexity

Time Complexity: O(n3) - Since each possible subarray is iterated over.
Space Complexity: O(n) - For storing distinct elements using a set.

Try this approach in the editor →

Approach 2: Sliding Window Technique

Using the Sliding Window Technique can improve efficiency by reducing repeated calculations of distinct counts. By maintaining a sliding window over the array, we can update the distinct count as the window slides, maintaining a dynamic result set of numbers.

This Python solution utilizes a sliding window technique with two pointers, left and right, representing the current window. A frequency array 'count' keeps track of element occurrences, and unique_count stores the number of distinct elements in the current window. After each iteration, the distinct count squared is added to total_sum.

Code

PythonC++

Python

C++

Complexity

Time Complexity: O(n)
Space Complexity: O(1) when considering the auxiliary space for counting and a fixed-size integer array.

Try this approach in the editor →

Complexity Comparison

ApproachComplexity
Brute Force Approach

Time Complexity: O(n3) - Since each possible subarray is iterated over.
Space Complexity: O(n) - For storing distinct elements using a set.

Sliding Window Technique

Time Complexity: O(n)
Space Complexity: O(1) when considering the auxiliary space for counting and a fixed-size integer array.

Detailed Complexity Analysis

ApproachTimeSpaceWhen to Use
Brute Force with SetO(n²)O(n)Small arrays or when demonstrating the base idea of counting distinct elements per subarray
Sliding Window + Fenwick/Segment TreeO(n log n)O(n)Large inputs where efficient range updates and prefix aggregation are required

Video Solution

Beautiful Segment tree lazy idea | Subarrays Distinct Element Sum of Squares II • Vivek Gupta • 4,650 views views

Watch 2 more video solutions →

Frequently Asked Questions

Is Subarrays Distinct Element Sum of Squares II easy or hard?
Subarrays Distinct Element Sum of Squares II is classified as a Hard problem. It requires understanding subarray contribution patterns, tracking last occurrences, and implementing advanced data structures such as Fenwick Trees or Segment Trees to achieve the optimal O(n log n) solution.
Subarrays Distinct Element Sum of Squares II Python/Java solution
Python and Java implementations typically use a Fenwick Tree or Segment Tree combined with a dictionary or array to track last occurrences. Each iteration updates contribution ranges and queries cumulative values to maintain the total sum of squared distinct counts.
How to solve Subarrays Distinct Element Sum of Squares II in O(n log n)?
Process the array from left to right while tracking the previous index of each value. When a value appears again, adjust the contribution range where it affects distinct counts. Use a Fenwick Tree or Segment Tree to apply range updates and compute prefix sums of squared distinct counts efficiently.
What is the best approach for Subarrays Distinct Element Sum of Squares II?
The most efficient solution uses a contribution-based technique with a Binary Indexed Tree (Fenwick Tree) or Segment Tree. While scanning the array, track the last occurrence of each element and update the range of subarrays whose distinct counts change. This reduces repeated work and achieves O(n log n) time complexity.
Is Subarrays Distinct Element Sum of Squares II asked at Google/Amazon/Meta?
Hard array and data structure problems involving segment trees or Fenwick trees frequently appear in interviews at companies like Google, Amazon, and Meta. This problem tests advanced subarray contribution techniques and efficient range updates, which are common interview patterns.
What data structure is used in Subarrays Distinct Element Sum of Squares II?
Efficient solutions rely on Binary Indexed Trees (Fenwick Trees) or Segment Trees to handle range updates and prefix queries. A hash map or array is also used to store the last occurrence index of each element while scanning the array.
What is the time complexity of Subarrays Distinct Element Sum of Squares II?
The brute force method runs in O(n^2) time because it enumerates every subarray and maintains a set of distinct elements. The optimized solution processes the array once and performs logarithmic range updates and queries using Fenwick or segment trees, resulting in O(n log n) time and O(n) space.

Ready to solve this problem?

Practice Subarrays Distinct Element Sum of Squares II with our built-in code editor and test cases.

Practice on FleetCode
Previous

Two Sum

Next

Median of Two Sorted Arrays