feat: add solutions to lc problem: No.2916

solution/2900-2999/2916.Subarrays Distinct Element Sum of Squares II/Solution.cpp

@@ -0,0 +1,86 @@
+class Node {
+public:
+    int l = 0, r = 0;
+    long long sum = 0, sumSq = 0;
+    int lazy = 0;
+};
+
+class SegmentTree {
+public:
+    SegmentTree(int n) {
+        tr.resize(n << 2);
+        build(1, 0, n - 1);
+    }
+
+    void modify(int u, int l, int r, int v) {
+        if (tr[u].l >= l && tr[u].r <= r) {
+            apply(u, v);
+            return;
+        }
+        pushdown(u);
+        int mid = (tr[u].l + tr[u].r) >> 1;
+        if (l <= mid) {
+            modify(u << 1, l, r, v);
+        }
+        if (r > mid) {
+            modify(u << 1 | 1, l, r, v);
+        }
+        pushup(u);
+    }
+
+    long long sumSq() const {
+        return tr[1].sumSq;
+    }
+
+private:
+    vector<Node> tr;
+
+    void build(int u, int l, int r) {
+        tr[u].l = l;
+        tr[u].r = r;
+        if (l == r) {
+            return;
+        }
+        int mid = (l + r) >> 1;
+        build(u << 1, l, mid);
+        build(u << 1 | 1, mid + 1, r);
+    }
+
+    void apply(int u, int v) {
+        long long len = tr[u].r - tr[u].l + 1LL;
+        long long add = v;
+        tr[u].sumSq += 2 * add * tr[u].sum + add * add * len;
+        tr[u].sum += add * len;
+        tr[u].lazy += v;
+    }
+
+    void pushdown(int u) {
+        if (tr[u].lazy) {
+            apply(u << 1, tr[u].lazy);
+            apply(u << 1 | 1, tr[u].lazy);
+            tr[u].lazy = 0;
+        }
+    }
+
+    void pushup(int u) {
+        tr[u].sum = tr[u << 1].sum + tr[u << 1 | 1].sum;
+        tr[u].sumSq = tr[u << 1].sumSq + tr[u << 1 | 1].sumSq;
+    }
+};
+
+class Solution {
+public:
+    int sumCounts(vector<int>& nums) {
+        int n = nums.size();
+        SegmentTree st(n);
+        unordered_map<int, int> last;
+        long long ans = 0;
+        for (int i = 0; i < n; ++i) {
+            int p = last.count(nums[i]) ? last[nums[i]] : -1;
+            st.modify(1, p + 1, i, 1);
+            last[nums[i]] = i;
+            ans = (ans + st.sumSq()) % 1000000007;
+        }
+        return (int) ans;
+    }
+};

solution/2900-2999/2916.Subarrays Distinct Element Sum of Squares II/Solution.cs

@@ -0,0 +1,85 @@
+public class Solution {
+    private const int Mod = 1000000007;
+
+    private class Node {
+        public int L, R;
+        public long Sum, SumSq;
+        public int Lazy;
+    }
+
+    private class SegmentTree {
+        private readonly Node[] tr;
+
+        public SegmentTree(int n) {
+            tr = new Node[n << 2];
+            for (int i = 0; i < tr.Length; i++) {
+                tr[i] = new Node();
+            }
+            Build(1, 0, n - 1);
+        }
+
+        public long SumSq => tr[1].SumSq;
+
+        private void Build(int u, int l, int r) {
+            tr[u].L = l;
+            tr[u].R = r;
+            if (l == r) {
+                return;
+            }
+            int mid = (l + r) >> 1;
+            Build(u << 1, l, mid);
+            Build(u << 1 | 1, mid + 1, r);
+        }
+
+        private void Apply(int u, int v) {
+            long len = tr[u].R - tr[u].L + 1L;
+            long add = v;
+            tr[u].SumSq += 2 * add * tr[u].Sum + add * add * len;
+            tr[u].Sum += add * len;
+            tr[u].Lazy += v;
+        }
+
+        private void PushDown(int u) {
+            if (tr[u].Lazy != 0) {
+                Apply(u << 1, tr[u].Lazy);
+                Apply(u << 1 | 1, tr[u].Lazy);
+                tr[u].Lazy = 0;
+            }
+        }
+
+        private void PushUp(int u) {
+            tr[u].Sum = tr[u << 1].Sum + tr[u << 1 | 1].Sum;
+            tr[u].SumSq = tr[u << 1].SumSq + tr[u << 1 | 1].SumSq;
+        }
+
+        public void Modify(int u, int l, int r, int v) {
+            if (tr[u].L >= l && tr[u].R <= r) {
+                Apply(u, v);
+                return;
+            }
+            PushDown(u);
+            int mid = (tr[u].L + tr[u].R) >> 1;
+            if (l <= mid) {
+                Modify(u << 1, l, r, v);
+            }
+            if (r > mid) {
+                Modify(u << 1 | 1, l, r, v);
+            }
+            PushUp(u);
+        }
+    }
+
+    public int SumCounts(int[] nums) {
+        int n = nums.Length;
+        SegmentTree st = new SegmentTree(n);
+        Dictionary<int, int> last = new Dictionary<int, int>();
+        long ans = 0;
+        for (int i = 0; i < n; i++) {
+            int p = last.TryGetValue(nums[i], out int idx) ? idx : -1;
+            st.Modify(1, p + 1, i, 1);
+            last[nums[i]] = i;
+            ans = (ans + st.SumSq) % Mod;
+        }
+        return (int) ans;
+    }
+}

solution/2900-2999/2916.Subarrays Distinct Element Sum of Squares II/Solution.go

@@ -0,0 +1,80 @@
+const mod = 1_000_000_007
+
+type node struct {
+	l, r       int
+	sum, sumSq int
+	lazy       int
+}
+
+type segmentTree struct {
+	tr []node
+}
+
+func newSegmentTree(n int) *segmentTree {
+	st := &segmentTree{tr: make([]node, n<<2)}
+	st.build(1, 0, n-1)
+	return st
+}
+
+func (st *segmentTree) build(u, l, r int) {
+	st.tr[u].l, st.tr[u].r = l, r
+	if l == r {
+		return
+	}
+	mid := (l + r) >> 1
+	st.build(u<<1, l, mid)
+	st.build(u<<1|1, mid+1, r)
+}
+
+func (st *segmentTree) apply(u, v int) {
+	length := st.tr[u].r - st.tr[u].l + 1
+	st.tr[u].sumSq += 2*v*st.tr[u].sum + v*v*length
+	st.tr[u].sum += v * length
+	st.tr[u].lazy += v
+}
+
+func (st *segmentTree) pushdown(u int) {
+	if st.tr[u].lazy != 0 {
+		st.apply(u<<1, st.tr[u].lazy)
+		st.apply(u<<1|1, st.tr[u].lazy)
+		st.tr[u].lazy = 0
+	}
+}
+
+func (st *segmentTree) pushup(u int) {
+	st.tr[u].sum = st.tr[u<<1].sum + st.tr[u<<1|1].sum
+	st.tr[u].sumSq = st.tr[u<<1].sumSq + st.tr[u<<1|1].sumSq
+}
+
+func (st *segmentTree) modify(u, l, r, v int) {
+	if st.tr[u].l >= l && st.tr[u].r <= r {
+		st.apply(u, v)
+		return
+	}
+	st.pushdown(u)
+	mid := (st.tr[u].l + st.tr[u].r) >> 1
+	if l <= mid {
+		st.modify(u<<1, l, r, v)
+	}
+	if r > mid {
+		st.modify(u<<1|1, l, r, v)
+	}
+	st.pushup(u)
+}
+
+func sumCounts(nums []int) int {
+	n := len(nums)
+	st := newSegmentTree(n)
+	last := map[int]int{}
+	ans := 0
+	for i, x := range nums {
+		p, ok := last[x]
+		if !ok {
+			p = -1
+		}
+		st.modify(1, p+1, i, 1)
+		last[x] = i
+		ans = (ans + st.tr[1].sumSq) % mod
+	}
+	return ans
+}

solution/2900-2999/2916.Subarrays Distinct Element Sum of Squares II/Solution.java

@@ -0,0 +1,83 @@
+class Solution {
+    private static final int MOD = 1000000007;
+
+    static class Node {
+        int l, r;
+        long sum, sumSq;
+        int lazy;
+    }
+
+    static class SegmentTree {
+        Node[] tr;
+
+        SegmentTree(int n) {
+            tr = new Node[n << 2];
+            for (int i = 0; i < tr.length; i++) {
+                tr[i] = new Node();
+            }
+            build(1, 0, n - 1);
+        }
+
+        void build(int u, int l, int r) {
+            tr[u].l = l;
+            tr[u].r = r;
+            if (l == r) {
+                return;
+            }
+            int mid = (l + r) >> 1;
+            build(u << 1, l, mid);
+            build(u << 1 | 1, mid + 1, r);
+        }
+
+        void apply(int u, int v) {
+            long len = tr[u].r - tr[u].l + 1L;
+            long add = v;
+            tr[u].sumSq += 2 * add * tr[u].sum + add * add * len;
+            tr[u].sum += add * len;
+            tr[u].lazy += v;
+        }
+
+        void pushdown(int u) {
+            if (tr[u].lazy != 0) {
+                apply(u << 1, tr[u].lazy);
+                apply(u << 1 | 1, tr[u].lazy);
+                tr[u].lazy = 0;
+            }
+        }
+
+        void pushup(int u) {
+            tr[u].sum = tr[u << 1].sum + tr[u << 1 | 1].sum;
+            tr[u].sumSq = tr[u << 1].sumSq + tr[u << 1 | 1].sumSq;
+        }
+
+        void modify(int u, int l, int r, int v) {
+            if (tr[u].l >= l && tr[u].r <= r) {
+                apply(u, v);
+                return;
+            }
+            pushdown(u);
+            int mid = (tr[u].l + tr[u].r) >> 1;
+            if (l <= mid) {
+                modify(u << 1, l, r, v);
+            }
+            if (r > mid) {
+                modify(u << 1 | 1, l, r, v);
+            }
+            pushup(u);
+        }
+    }
+
+    public int sumCounts(int[] nums) {
+        int n = nums.length;
+        SegmentTree st = new SegmentTree(n);
+        Map<Integer, Integer> last = new HashMap<>();
+        long ans = 0;
+        for (int i = 0; i < n; i++) {
+            int p = last.getOrDefault(nums[i], -1);
+            st.modify(1, p + 1, i, 1);
+            last.put(nums[i], i);
+            ans = (ans + st.tr[1].sumSq) % MOD;
+        }
+        return (int) ans;
+    }
+}

solution/2900-2999/2916.Subarrays Distinct Element Sum of Squares II/Solution.py

@@ -0,0 +1,61 @@
+MOD = 10**9 + 7
+
+
+class Solution:
+    def sumCounts(self, nums: List[int]) -> int:
+        n = len(nums)
+
+        class Node:
+            __slots__ = ("l", "r", "sum", "sum_sq", "lazy")
+
+            def __init__(self):
+                self.l = self.r = 0
+                self.sum = self.sum_sq = self.lazy = 0
+
+        tr = [Node() for _ in range(n << 2)]
+
+        def build(u: int, l: int, r: int) -> None:
+            tr[u].l, tr[u].r = l, r
+            if l == r:
+                return
+            mid = (l + r) >> 1
+            build(u << 1, l, mid)
+            build(u << 1 | 1, mid + 1, r)
+
+        def apply(u: int, v: int) -> None:
+            length = tr[u].r - tr[u].l + 1
+            tr[u].sum_sq += 2 * v * tr[u].sum + v * v * length
+            tr[u].sum += v * length
+            tr[u].lazy += v
+
+        def pushdown(u: int) -> None:
+            if tr[u].lazy:
+                apply(u << 1, tr[u].lazy)
+                apply(u << 1 | 1, tr[u].lazy)
+                tr[u].lazy = 0
+
+        def pushup(u: int) -> None:
+            tr[u].sum = tr[u << 1].sum + tr[u << 1 | 1].sum
+            tr[u].sum_sq = tr[u << 1].sum_sq + tr[u << 1 | 1].sum_sq
+
+        def modify(u: int, l: int, r: int, v: int) -> None:
+            if tr[u].l >= l and tr[u].r <= r:
+                apply(u, v)
+                return
+            pushdown(u)
+            mid = (tr[u].l + tr[u].r) >> 1
+            if l <= mid:
+                modify(u << 1, l, r, v)
+            if r > mid:
+                modify(u << 1 | 1, l, r, v)
+            pushup(u)
+
+        build(1, 0, n - 1)
+        last = {}
+        ans = 0
+        for i, x in enumerate(nums):
+            p = last.get(x, -1)
+            modify(1, p + 1, i, 1)
+            last[x] = i
+            ans = (ans + tr[1].sum_sq) % MOD
+        return ans