2916. Subarrays Distinct Element Sum of Squares II

/**
 * Problem: LeetCode 2916 - Subarrays Distinct Element Sum of Squares II
 * Difficulty: Hard
 * Language: Java
 * Author: CodeByTushu
 */
import java.util.Arrays;

/**
 * Problem: LeetCode 2916 - Subarrays Distinct Element Sum of Squares II
 * Difficulty: Hard
 * Language: Java
 * Author: CodeByTushu
 */
class Solution {
    private static final long MOD = 1_000_000_007L;
    private long[] sum;
    private long[] sq;
    private long[] lazy;

    public int sumCounts(int[] nums) {
        int n = nums.length;
        sum = new long[4 * n];
        sq = new long[4 * n];
        lazy = new long[4 * n];

        int[] last = new int[100001];
        Arrays.fill(last, -1);

        long ans = 0;
        for (int j = 0; j < n; j++) {
            int x = nums[j];
            update(1, 0, n - 1, last[x] + 1, j, 1);
            last[x] = j;
            ans = (ans + sq[1]) % MOD;
        }

        return (int) ans;
    }

    private void apply(int node, int len, long v) {
        sq[node] = (sq[node] + 2 * v % MOD * sum[node] % MOD + v * v % MOD * len) % MOD;
        sum[node] = (sum[node] + v * len) % MOD;
        lazy[node] = (lazy[node] + v) % MOD;
    }

    private void push(int node, int l, int r) {
        if (lazy[node] != 0) {
            int mid = (l + r) / 2;
            apply(node * 2, mid - l + 1, lazy[node]);
            apply(node * 2 + 1, r - mid, lazy[node]);
            lazy[node] = 0;
        }
    }

    private void update(int node, int l, int r, int ql, int qr, long v) {
        if (qr < l || r < ql) return;
        if (ql <= l && r <= qr) {
            apply(node, r - l + 1, v);
            return;
        }
        push(node, l, r);
        int mid = (l + r) / 2;
        update(node * 2, l, mid, ql, qr, v);
        update(node * 2 + 1, mid + 1, r, ql, qr, v);
        sum[node] = (sum[node * 2] + sum[node * 2 + 1]) % MOD;
        sq[node] = (sq[node * 2] + sq[node * 2 + 1]) % MOD;
    }
}

"""
Problem: LeetCode 2916 - Subarrays Distinct Element Sum of Squares II
Difficulty: Hard
Language: Python
Author: CodeByTushu
"""
from typing import List

class Solution:
    def sumCounts(self, nums: List[int]) -> int:
        MOD = 1_000_000_007
        n = len(nums)

        tree_sum = [0] * (4 * n)
        tree_sq = [0] * (4 * n)
        lazy = [0] * (4 * n)

        def apply(node: int, length: int, v: int) -> None:
            tree_sq[node] = (tree_sq[node] + 2 * v * tree_sum[node] + v * v * length) % MOD
            tree_sum[node] = (tree_sum[node] + v * length) % MOD
            lazy[node] = (lazy[node] + v) % MOD

        def push(node: int, l: int, r: int) -> None:
            if lazy[node]:
                mid = (l + r) // 2
                apply(node * 2, mid - l + 1, lazy[node])
                apply(node * 2 + 1, r - mid, lazy[node])
                lazy[node] = 0

        def update(node: int, l: int, r: int, ql: int, qr: int, v: int) -> None:
            if qr < l or r < ql:
                return
            if ql <= l and r <= qr:
                apply(node, r - l + 1, v)
                return
            push(node, l, r)
            mid = (l + r) // 2
            update(node * 2, l, mid, ql, qr, v)
            update(node * 2 + 1, mid + 1, r, ql, qr, v)
            tree_sum[node] = (tree_sum[node * 2] + tree_sum[node * 2 + 1]) % MOD
            tree_sq[node] = (tree_sq[node * 2] + tree_sq[node * 2 + 1]) % MOD

        last = {}
        ans = 0

        for j, x in enumerate(nums):
            prev = last.get(x, -1)
            update(1, 0, n - 1, prev + 1, j, 1)
            last[x] = j
            ans = (ans + tree_sq[1]) % MOD

        return ans

/**
 * Problem: LeetCode 2916 - Subarrays Distinct Element Sum of Squares II
 * Difficulty: Hard
 * Language: C++
 * Author: CodeByTushu
 */
#include <vector>

using namespace std;

/**
 * Problem: LeetCode 2916 - Subarrays Distinct Element Sum of Squares II
 * Difficulty: Hard
 * Language: C++
 * Author: CodeByTushu
 */
class Solution {
private:
    static constexpr long long MOD = 1000000007LL;
    vector<long long> tree_sum;
    vector<long long> tree_sq;
    vector<long long> lazy;

    void apply(int node, int len, long long v) {
        tree_sq[node] = (tree_sq[node] + 2LL * v % MOD * tree_sum[node] % MOD + v * v % MOD * len) % MOD;
        tree_sum[node] = (tree_sum[node] + v * len) % MOD;
        lazy[node] = (lazy[node] + v) % MOD;
    }

    void push(int node, int l, int r) {
        if (lazy[node] != 0) {
            int mid = (l + r) / 2;
            apply(node * 2, mid - l + 1, lazy[node]);
            apply(node * 2 + 1, r - mid, lazy[node]);
            lazy[node] = 0;
        }
    }

    void update(int node, int l, int r, int ql, int qr, long long v) {
        if (qr < l || r < ql) return;
        if (ql <= l && r <= qr) {
            apply(node, r - l + 1, v);
            return;
        }
        push(node, l, r);
        int mid = (l + r) / 2;
        update(node * 2, l, mid, ql, qr, v);
        update(node * 2 + 1, mid + 1, r, ql, qr, v);
        tree_sum[node] = (tree_sum[node * 2] + tree_sum[node * 2 + 1]) % MOD;
        tree_sq[node] = (tree_sq[node * 2] + tree_sq[node * 2 + 1]) % MOD;
    }

public:
    int sumCounts(vector<int>& nums) {
        int n = nums.size();
        tree_sum.assign(4 * n, 0);
        tree_sq.assign(4 * n, 0);
        lazy.assign(4 * n, 0);

        int max_val = 0;
        for (int x : nums) {
            if (x > max_val) max_val = x;
        }
        vector<int> last(max_val + 1, -1);

        long long ans = 0;
        for (int j = 0; j < n; j++) {
            int x = nums[j];
            update(1, 0, n - 1, last[x] + 1, j, 1);
            last[x] = j;
            ans = (ans + tree_sq[1]) % MOD;
        }

        return ans;
    }
};

/**
 * Problem: LeetCode 2916 - Subarrays Distinct Element Sum of Squares II
 * Difficulty: Hard
 * Language: JavaScript
 * Author: CodeByTushu
 */
var sumCounts = function(nums) {
    const MOD = 1000000007n;
    const n = nums.length;

    const sum = new BigInt64Array(4 * n);
    const sq = new BigInt64Array(4 * n);
    const lazy = new BigInt64Array(4 * n);

    function apply(node, len, v) {
        const bigLen = BigInt(len);
        sq[node] = (sq[node] + 2n * v * sum[node] + v * v * bigLen) % MOD;
        sum[node] = (sum[node] + v * bigLen) % MOD;
        lazy[node] = (lazy[node] + v) % MOD;
    }

    function push(node, l, r) {
        if (lazy[node] !== 0n) {
            const mid = (l + r) >> 1;
            apply(node * 2, mid - l + 1, lazy[node]);
            apply(node * 2 + 1, r - mid, lazy[node]);
            lazy[node] = 0n;
        }
    }

    function update(node, l, r, ql, qr, v) {
        if (qr < l || r < ql) return;
        if (ql <= l && r <= qr) {
            apply(node, r - l + 1, v);
            return;
        }
        push(node, l, r);
        const mid = (l + r) >> 1;
        update(node * 2, l, mid, ql, qr, v);
        update(node * 2 + 1, mid + 1, r, ql, qr, v);
        sum[node] = (sum[node * 2] + sum[node * 2 + 1]) % MOD;
        sq[node] = (sq[node * 2] + sq[node * 2 + 1]) % MOD;
    }

    let maxVal = 0;
    for (let i = 0; i < n; i++) {
        if (nums[i] > maxVal) maxVal = nums[i];
    }
    const last = new Int32Array(maxVal + 1);
    last.fill(-1);

    let ans = 0n;
    for (let j = 0; j < n; j++) {
        const x = nums[j];
        update(1, 0, n - 1, last[x] + 1, j, 1n);
        last[x] = j;
        ans = (ans + sq[1]) % MOD;
    }

    return Number(ans);
};