2916. Subarrays Distinct Element Sum of Squares II

HardView on LeetCode

Array Hash Table Binary Indexed Tree

Time: O(n log n)

Space: O(n + V)

Problem Overview

This is the same counting problem as the easy version, but n reaches 1e5, so every subarray cannot be scanned.

Intuition

This is the same counting problem as the easy version, but n reaches 1e5, so every subarray cannot be scanned. Fix the right end r and let c[l] be how many distinct values sit in nums[l..r]. Moving r forward adds nums[r], which raises c[l] by one only for left ends after the previous place that value appeared. Raising a count from c to c + 1 raises its square by 2c + 1, so the sum of squares over all left ends grows by twice the sum of the old counts on that range plus the range length. Range add and range sum are enough, and a pair of Fenwick trees answers both in logarithmic time per step.

Algorithm

1Keep two Fenwick trees that together support adding a value on a range and summing counts on a range.
2Store the last index of each value in a fixed array sized to the value limit.
3For each right end r, let l be one past the previous occurrence of nums[r], or 1 if it is new.
4Add twice the current sum of counts from l through r, plus the length of that range, to a running total of squares.
5Add 1 on the range [l, r], record r as the last place nums[r] appeared, and add the running total to the answer modulo 1e9+7.

Example Walkthrough

Input: nums = [1,2,1]

1.r = 1: only left end 1, so the total of squares becomes 1.
2.r = 2: value 2 is new, so both left ends gain one. The old counts sum to 1, so the total grows by 4 and becomes 5.
3.r = 3: value 1 was last at index 1, so only left ends 2 and 3 gain one. The total grows by 4 and becomes 9.
4.The per-right-end totals 1, 5, and 9 add up to 15.

Output: 15

Common Pitfalls

•Only left ends after the previous occurrence gain a new distinct value. Updating every left end overcounts repeats.
•Read the range sum before adding 1, because the 2c + 1 formula uses the old counts.
•The running sum of squares can reach about n cubed over 3, so keep it in a 64-bit integer and reduce only the answer.
•A lazy segment tree that stores sums of squares also works, but two Fenwick trees use less memory and less constant time.

2916.cs

// Approach: Sweep the right end r and let c[l] be the distinct count of
// nums[l..r]. Adding nums[r] raises c[l] by 1 for every l after the previous
// occurrence of nums[r]. Since (c + 1)^2 - c^2 = 2c + 1, the sum of squares
// over all l grows by 2 * (sum of c over that range) + (range length). So only
// range add and range sum are needed, which two Fenwick trees answer
// iteratively. The running sum of squares is the total for subarrays ending at
// r, and those totals add up to the answer.
// Time: O(n log n) Space: O(n + V), where V is the value range

public class Solution
{
    private long[] b1;
    private long[] b2;
    private int size;

    public int SumCounts(int[] nums)
    {
        const int kMod = 1_000_000_007;
        int n = nums.Length;
        size = n;
        b1 = new long[n + 1];
        b2 = new long[n + 1];
        int[] lastSeen = new int[100_001];

        long squares = 0;
        long ans = 0;
        for (int r = 1; r <= n; r++)
        {
            int x = nums[r - 1];
            int l = lastSeen[x] + 1;
            lastSeen[x] = r;
            squares += 2 * (Prefix(r) - Prefix(l - 1)) + (r - l + 1);
            Add(l, 1);
            Add(r + 1, -1);
            ans = (ans + squares) % kMod;
        }
        return (int)ans;
    }

    private void Add(int i, long x)
    {
        if (i > size)
            return;
        long shifted = x * (i - 1);
        for (; i <= size; i += i & -i)
        {
            b1[i] += x;
            b2[i] += shifted;
        }
    }

    private long Prefix(int i)
    {
        long s1 = 0, s2 = 0;
        for (int j = i; j > 0; j -= j & -j)
        {
            s1 += b1[j];
            s2 += b2[j];
        }
        return s1 * i - s2;
    }
}

Was this solution helpful?