DDSA Solutions

2913. Subarrays Distinct Element Sum of Squares I

EasyView on LeetCode
ArrayHash TableBinary Indexed Tree
Time: O(n log n)
Space: O(n + V)

Problem Overview

Fix the right end r of a subarray and track, for every left end l, how many distinct values nums[l..r] holds.

Intuition

Fix the right end r of a subarray and track, for every left end l, how many distinct values nums[l..r] holds. Moving r one step adds nums[r]. That raises the count by one only for left ends after the previous place nums[r] appeared, because earlier left ends already contain it. Raising a count c to c + 1 raises its square by 2c + 1. So the total of squares for all subarrays ending at r grows by twice the sum of the old counts in that range plus the range length. Range add and range sum are exactly what a pair of Fenwick trees supports.

Algorithm

  1. 1Keep two Fenwick trees that together support adding a value to a range and summing a range of counts.
  2. 2For each right end r, look up the last index where nums[r] appeared. The affected left ends run from just after it through r.
  3. 3Add twice the current sum of counts over that range, plus its length, to a running total of squares.
  4. 4Add 1 to every count in that range and record r as the last place nums[r] appeared.
  5. 5Add the running total to the answer, which now covers every subarray ending at r.

Example Walkthrough

Input: nums = [1,2,1]

  1. 1.r = 0: the range is just index 0, so the total of squares becomes 1.
  2. 2.r = 1: value 2 is new, so both left ends gain one. The old counts sum to 1, so the total grows by 2 times 1 plus 2 and becomes 5.
  3. 3.r = 2: value 1 was last at index 0, so only left ends 1 and 2 gain one. Their old counts sum to 1, so the total grows by 4 and becomes 9.
  4. 4.The totals 1, 5, and 9 add up to the answer.

Output: 15

Common Pitfalls

  • •Only left ends after the previous occurrence gain a new distinct value. Updating every left end overcounts repeats.
  • •Read the sum of counts before adding 1, because the 2c + 1 formula uses the old counts.
  • •The running total can reach about n cubed over 3, so keep it in a 64-bit integer.
  • •Checking every subarray with a set takes quadratic time. That passes here but does not scale to the larger version of the problem.
2913.cs
C#
// 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 number of distinct values

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

    public int SumCounts(IList<int> nums)
    {
        const int kMod = 1_000_000_007;
        int n = nums.Count;
        size = n;
        b1 = new long[n + 1];
        b2 = new long[n + 1];
        var lastSeen = new Dictionary<int, int>();

        long squares = 0;
        long ans = 0;
        for (int r = 1; r <= n; r++)
        {
            int l = lastSeen.TryGetValue(nums[r - 1], out int prev) ? prev + 1 : 1;
            lastSeen[nums[r - 1]] = 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)
    {
        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?

Related Problems