2916. Subarrays Distinct Element Sum of Squares II LeetCode Solution

Last updated: March 5, 2025

In this guide, you will get 2916. Subarrays Distinct Element Sum of Squares II LeetCode Solution with the best time and space complexity. The solution to Subarrays Distinct Element Sum of Squares II problem is provided in various programming languages like C++, Java, and Python. This will be helpful for you if you are preparing for placements, hackathons, interviews, or practice purposes. The solutions provided here are very easy to follow and include detailed explanations.

Problem Statement
Complexity Analysis
Subarrays Distinct Element Sum of Squares II solution in C++
Subarrays Distinct Element Sum of Squares II solution in Java
Subarrays Distinct Element Sum of Squares II solution in Python
Additional Resources

2916. Subarrays Distinct Element Sum of Squares II LeetCode Solution image

Problem Statement of Subarrays Distinct Element Sum of Squares II

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

Complexity Analysis

Time Complexity: O(n\log n)
Space Complexity: O(n)

2916. Subarrays Distinct Element Sum of Squares II LeetCode Solution in C++

class SegmentTree {
 public:
  SegmentTree(int n, int kMod)
      : n(n), kMod(kMod), lazy(4 * n), sums(4 * n), squaredSums(4 * n) {}

  void updateRange(int l, int r) {
    return updateRange(/*i=*/0, /*start=*/0, /*end=*/n - 1, l, r);
  }

  void propagate(int i, int l, int r) {
    const int gap = r - l + 1;
    // (a + L)^2 + (b + L)^2 + (c + L)^2, where L = lazy[i]
    // a^2 + b^2 + c^2 + 2 * L (a + b + c) + L^2 * gap, where gap = 3
    squaredSums[i] += 2 * lazy[i] * sums[i] + lazy[i] * lazy[i] * gap;
    squaredSums[i] %= kMod;
    sums[i] += lazy[i] * gap;
    sums[i] %= kMod;
    if (l < r) {
      lazy[i * 2 + 1] += lazy[i];
      lazy[i * 2 + 2] += lazy[i];
    }
    lazy[i] = 0;
  }

  int getTreeSquaredSums() {
    return squaredSums[0];
  }

 private:
  const int kMod;
  const int n;
  vector<long> lazy;
  vector<long> sums;
  vector<long> squaredSums;

  void updateRange(int i, int start, int end, int l, int r) {
    if (lazy[i] > 0)
      propagate(i, start, end);
    if (end < l || start > r)
      return;
    if (start >= l && end <= r) {
      lazy[i] = 1;
      propagate(i, start, end);
      return;
    }
    const int mid = (start + end) / 2;
    updateRange(i * 2 + 1, start, mid, l, r);
    updateRange(i * 2 + 2, mid + 1, end, l, r);
    sums[i] = (sums[i * 2 + 1] + sums[i * 2 + 2]) % kMod;
    squaredSums[i] = (squaredSums[i * 2 + 1] + squaredSums[i * 2 + 2]) % kMod;
  }
};

class Solution {
 public:
  int sumCounts(vector<int>& nums) {
    constexpr int kMod = 1'000'000'007;
    const int n = nums.size();
    int ans = 0;
    unordered_map<int, int> lastSeen;
    SegmentTree tree(n, kMod);

    for (int r = 0; r < n; ++r) {
      const int l = lastSeen.contains(nums[r]) ? lastSeen[nums[r]] + 1 : 0;
      tree.updateRange(l, r);
      lastSeen[nums[r]] = r;
      ans = (ans + tree.getTreeSquaredSums()) % kMod;
    }

    return ans;
  }
};
/* code provided by PROGIEZ */

2916. Subarrays Distinct Element Sum of Squares II LeetCode Solution in Java

N/A
// code provided by PROGIEZ

2916. Subarrays Distinct Element Sum of Squares II LeetCode Solution in Python

N/A
 # code by PROGIEZ

Additional Resources

Happy Coding! Keep following PROGIEZ for more updates and solutions.

Table of Contents