Subarrays Distinct Element Sum of Squares II

Subarrays Distinct Element Sum of Squares II - Problem

Imagine you're analyzing the diversity of elements in every possible contiguous section of an array. Your task is to measure this diversity in a unique way!

Given a 0-indexed integer array nums, you need to:

Find all possible subarrays - every contiguous sequence of elements
Count distinct elements in each subarray
Square these counts and sum them all up

For example, if a subarray [1, 2, 1] has 2 distinct elements (1 and 2), it contributes 2² = 4 to our final answer.

Return the sum modulo 10⁹ + 7 since the result can be extremely large.

Challenge: With arrays up to 10⁵ elements, a naive approach won't suffice - you'll need advanced data structures like Binary Indexed Trees or Segment Trees to efficiently track contribution changes as subarrays expand.

Input & Output

example_1.py — Simple Array

$ Input: nums = [1, 2, 1]

› Output: 15

💡 Note: Subarray [1] has 1 distinct element: 1² = 1. Subarray [2] has 1 distinct element: 1² = 1. Subarray [1] has 1 distinct element: 1² = 1. Subarray [1,2] has 2 distinct elements: 2² = 4. Subarray [2,1] has 2 distinct elements: 2² = 4. Subarray [1,2,1] has 2 distinct elements: 2² = 4. Total: 1+1+1+4+4+4 = 15

example_2.py — All Same Elements

$ Input: nums = [2, 2, 2]

› Output: 6

💡 Note: Every subarray has exactly 1 distinct element (the value 2). There are 6 subarrays total: [2], [2], [2], [2,2], [2,2], [2,2,2]. Each contributes 1² = 1 to the sum. Total: 6 × 1 = 6

example_3.py — All Different Elements

$ Input: nums = [1, 2, 3]

› Output: 20

💡 Note: Subarrays: [1]=1², [2]=1², [3]=1², [1,2]=2², [2,3]=2², [1,2,3]=3². Total: 1+1+1+4+4+9 = 20

Constraints

1 ≤ nums.length ≤ 10⁵
1 ≤ nums[i] ≤ 10⁵
Result must be returned modulo 10⁹ + 7

Visualization

Tap to expand

Understanding the Visualization

Setup tracking system

Initialize BIT to efficiently track contribution changes

Process each room

As visitors enter each room, update diversity contributions

Handle revisits

When artwork repeats, adjust previous contributions

Calculate final score

Sum all squared diversity scores across all paths

Key Takeaway

🎯 Key Insight: Instead of recalculating diversity for every path, we incrementally track how each new room affects all existing paths, using BIT for efficient updates when artworks repeat.

Asked in

G Google 23 f Meta 18 a Amazon 15 ⊞ Microsoft 12

The optimal solution uses a Binary Indexed Tree (Fenwick Tree) to efficiently track and update distinct element contributions as we process each position. By maintaining the last occurrence of each element and using the BIT for O(log n) updates and queries, we achieve O(n log n) time complexity instead of the naive O(n³) approach.

Common Approaches

Approach	Time	Space	Notes
✓ Binary Indexed Tree (Optimal)	O(n log n)	O(n)	Use BIT to efficiently track distinct count contributions dynamically
Brute Force (All Subarrays)	O(n³)	O(n)	Generate all subarrays and count distinct elements in each

Binary Indexed Tree (Optimal) — Algorithm Steps

Initialize Binary Indexed Tree and last occurrence map
For each position i, process element nums[i]
Update BIT with contribution changes from previous occurrences
Add current element's contribution to all subarrays ending at i
Accumulate the squared distinct counts efficiently

Visualization

Tap to expand

Step-by-Step Walkthrough

Initialize BIT

Set up Binary Indexed Tree to track contributions

Process each element

Update contributions based on last occurrence

Query efficiently

Use BIT to get sum of contributions in O(log n)

Code -

solution.c — C

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define MOD 1000000007
#define MAXN 100005

typedef struct {
    int size;
    long long tree[MAXN];
} BIT;

void bit_init(BIT* bit, int size) {
    bit->size = size;
    memset(bit->tree, 0, sizeof(bit->tree));
}

void bit_update(BIT* bit, int i, long long delta) {
    i++;
    while (i <= bit->size) {
        bit->tree[i] = (bit->tree[i] + delta + MOD) % MOD;
        i += i & (-i);
    }
}

long long bit_query(BIT* bit, int i) {
    i++;
    long long result = 0;
    while (i > 0) {
        result = (result + bit->tree[i]) % MOD;
        i -= i & (-i);
    }
    return result;
}

int sumCounts(int* nums, int n) {
    BIT bit;
    bit_init(&bit, n);
    
    int lastPos[MAXN];
    memset(lastPos, -1, sizeof(lastPos));
    
    long long result = 0;
    
    for (int i = 0; i < n; i++) {
        if (lastPos[nums[i]] != -1) {
            int prevPos = lastPos[nums[i]];
            long long oldContrib = prevPos > 0 ? 
                (bit_query(&bit, prevPos) - bit_query(&bit, prevPos - 1) + MOD) % MOD : 
                bit_query(&bit, 0);
            bit_update(&bit, prevPos, -oldContrib);
        }
        
        long long contribution = i > 0 ? 
            (2 * bit_query(&bit, i - 1) + 1) % MOD : 1;
        bit_update(&bit, i, contribution);
        lastPos[nums[i]] = i;
        
        result = (result + bit_query(&bit, i)) % MOD;
    }
    
    return result;
}

int main() {
    char line[1000000];
    fgets(line, sizeof(line), stdin);
    
    int nums[MAXN];
    int n = 0;
    char* token = strtok(line, "[],");
    while (token != NULL) {
        nums[n++] = atoi(token);
        token = strtok(NULL, "[],");
    }
    
    printf("%d\n", sumCounts(nums, n));
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n log n)

Single pass through array with O(log n) BIT operations per element

⚡ Linearithmic

Space Complexity

O(n)

BIT structure and hash map for last occurrences

⚡ Linearithmic Space

26.4K Views

Medium Frequency

~35 min Avg. Time

847 Likes

Ln 1, Col 1

Smart Actions

💡 Explanation

AI Ready

💡 Suggestion Tab to accept Esc to dismiss

// Output will appear here after running code

Code Editor Closed

Click the red button to reopen

Input & Output

Constraints

Visualization

Related Problems

Common Approaches

Binary Indexed Tree (Optimal) — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler