#2916Easy~30 min

Subarrays Distinct Element Sum of Squares II

Time O(nlogn) · Space O(n) · Official statement on LeetCode

dynamic programming dynamic programming segment tree fenwick tree sorted list

Solutions

// Time:  O(nlogn)
// Space: O(n)

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;

// bit, fenwick tree, ordered set, math
class Solution {
public:
    int sumCounts(vector<int>& nums) {
        static const int MOD = 1e9 + 7;

        using ordered_set = tree<int, null_type, less<int>, rb_tree_tag, tree_order_statistics_node_update>;
        ordered_set os;
        unordered_map<int, vector<int>> idxs;
        for (int i = size(nums) - 1; i >= 0; --i) {
            idxs[nums[i]].emplace_back(i);
        }
        for (const auto& [_, v] : idxs) {
            os.insert(v.back());
        }
        int result = 0;
        int accu = ((static_cast<int64_t>(size(nums)) * size(os)) % MOD) * size(os) % MOD;
        for (int i = 0; i < size(os); ++i) {
            accu = ((accu - (static_cast<int64_t>(2 * i + 1) * *os.find_by_order(i) % MOD)) % MOD + MOD) % MOD;
        }
        BIT bit(size(nums));
        for (const auto& x : os) {
            bit.add(x, x);
        }

        const auto& update = [&](int x, int accu, int d) {
            const int i = os.order_of_key(idxs[x].back());
            accu = ((accu + d * (static_cast<int64_t>(size(nums)) * (2 * static_cast<int>(size(os)) - 1)
                                 - static_cast<int64_t>(2 * i + 1) * idxs[x].back()
                                 - 2ll * (bit.query(size(nums) - 1) - bit.query(idxs[x].back())))) % MOD + MOD) % MOD;
            bit.add(idxs[x].back(), d * idxs[x].back());
            return accu;
        };

        for (const auto& x : nums) {
            result = (result + accu) % MOD;  // accu = sum(count(i, k) for k in range(i, len(nums)))
            accu = update(x, accu, -1);
            os.erase(idxs[x].back());
            idxs[x].pop_back();
            if (empty(idxs[x])) {
                continue;
            }
            os.insert(idxs[x].back());
            accu = update(x, accu, +1);
        }
        assert(accu == 0);
        return result;
    }

private:
    class BIT {
    public:
        BIT(int n) : bit_(n + 1) {  // 0-indexed
        }
        
        void add(int i, int val) {
            ++i;
            for (; i < size(bit_); i += lower_bit(i)) {
                bit_[i] = (bit_[i] + val) % MOD;
            }
        }

        int query(int i) const {
            ++i;
            int total = 0;
            for (; i > 0; i -= lower_bit(i)) {
                total = (total + bit_[i]) % MOD;
            }
            return total;
        }
    
    private:
        int lower_bit(int i) const {
            return i & -i;
        }
        
        vector<int> bit_;
        static const int MOD = 1e9 + 7;
    };
};

// Time:  O(nlogn)
// Space: O(n)
// dp, segment tree, math
class Solution2 {
public:
    int sumCounts(vector<int>& nums) {
        static const int MOD = 1e9 + 7;

        int result = 0, accu = 0;
        unordered_map<int, int> lookup;
        SegmentTree st(size(nums));
        for (int i = 0; i < size(nums); ++i) {
            const int j = lookup.count(nums[i]) ? lookup[nums[i]] : -1;
            // sum(count(k, i)^2 for k in range(i+1)) - sum(count(k, i-1)^2 for k in range(i))
            // = sum(2*count(k, i-1)+1 for k in range(j+1, i+1))
            // = (i-j) + sum(2*count(k, i-1) for k in range(j+1, i+1))
            accu = (accu + ((i - j) + 2ll * st.query(j + 1, i))) % MOD;
            result = (result + accu) % MOD;
            st.update(j + 1, i, 1);  // count(k, i) = count(k, i-1)+(1 if k >= j+1 else 0) for k in range(i+1)
            lookup[nums[i]] = i;
        }
        return result;
    }

private:
    class SegmentTree {
    private:
        static const int MOD = 1e9 + 7;

    public:
        explicit SegmentTree(
            int N)
         :  base_(N > 1 ? 1 << (__lg(N - 1) + 1) : 1),
            lazy_(base_),
            tree_(N > 1 ? 1 << (__lg(N - 1) + 2) : 2),
            count_(size(tree_), 1) {

            for (int i = base_ - 1; i >= 1; --i) { // added
                count_[i] = count_[i << 1] + count_[(i << 1) + 1];
            }
        }

        void update(int L, int R, const int val) {
            L += base_;
            R += base_;
            // push(L);  // enable if range assignment
            // push(R);  // enable if range assignment
            int L0 = L, R0 = R;
            for (; L <= R; L >>= 1, R >>= 1) {
                if ((L & 1) == 1) {
                    apply(L++, val);
                }
                if ((R & 1) == 0) {
                    apply(R--, val);
                }
            }
            pull(L0);
            pull(R0);
        }

        int query(int L, int R) {
            if (L > R) {
                return 0;
            }
            L += base_;
            R += base_;
            push(L);
            push(R);
            int left = 0, right = 0;
            for (; L <= R; L >>= 1, R >>= 1) {
                if ((L & 1) == 1) {
                    left = (left + tree_[L++]) % MOD;
                }
                if ((R & 1) == 0) {
                    right = (tree_[R--] + right) % MOD;
                }
            }
            return (left + right) % MOD;
        }

    private:
        void apply(int x, const int val) {
            tree_[x] = (tree_[x] + static_cast<int64_t>(val) * count_[x]) % MOD;  // modified
            if (x < base_) {
                lazy_[x] = (lazy_[x] + val) % MOD;
            }
        }

        void pull(int x) {
            while (x > 1) {
                x >>= 1;
                tree_[x] = (tree_[x << 1] + tree_[(x << 1) + 1]) % MOD;
                if (lazy_[x]) {
                    tree_[x] = (tree_[x] + static_cast<int64_t>(lazy_[x]) * count_[x]) % MOD;  // modified
                }
            }
        }

        void push(int x) {
            for (int h = __lg(x) - 1; h > 0; --h) {
                int y = x >> h;
                if (lazy_[y]) {
                    apply(y << 1, lazy_[y]);
                    apply((y << 1) + 1, lazy_[y]);
                    lazy_[y] = 0;
                }
            }
        }

        int base_;
        vector<int> tree_;
        vector<int> lazy_;
        vector<int> count_;  // added
    };
};

// Time:  O(nlogn)
// Space: O(n)
// dp, segment tree, math
class Solution3 {
public:
    int sumCounts(vector<int>& nums) {
        static const int MOD = 1e9 + 7;

        const auto& sum = [&] (const auto& x, const auto& y) {
            return (x + y) % MOD;
        };

        int result = 0, accu = 0;
        unordered_map<int, int> lookup;
        SegmentTree<int> st(size(nums), sum, sum);
        for (int i = 0; i < size(nums); ++i) {
            const int j = lookup.count(nums[i]) ? lookup[nums[i]] : -1;
            // sum(count(k, i)^2 for k in range(i+1)) - sum(count(k, i-1)^2 for k in range(i))
            // = sum(2*count(k, i-1)+1 for k in range(j+1, i+1))
            // = (i-j) + sum(2*count(k, i-1) for k in range(j+1, i+1))
            accu = (accu + ((i - j) + 2ll * st.query(j + 1, i))) % MOD;
            result = (result + accu) % MOD;
            st.update(j + 1, i, 1);  // count(k, i) = count(k, i-1)+(1 if k >= j+1 else 0) for k in range(i+1)
            lookup[nums[i]] = i;
        }
        return result;
    }

private:
    template <typename T>
    class SegmentTree {
    private:
        static const int MOD = 1e9 + 7;

    public:
        explicit SegmentTree(
            int N,
            const function<T(const T&, const T&)>& query_fn,
            const function<T(const T&, const T&)>& update_fn)
         :  base_(N > 1 ? 1 << (__lg(N - 1) + 1) : 1),
            lazy_(base_),
            tree_(N > 1 ? 1 << (__lg(N - 1) + 2) : 2),
            count_(size(tree_), 1),
            query_fn_(query_fn),
            update_fn_(update_fn) {

            for (int i = base_ - 1; i >= 1; --i) { // added
                count_[i] = count_[i << 1] + count_[(i << 1) + 1];
            }
        }

        void update(int L, int R, const T& val) {
            L += base_;
            R += base_;
            // push(L);  // enable if range assignment
            // push(R);  // enable if range assignment
            int L0 = L, R0 = R;
            for (; L <= R; L >>= 1, R >>= 1) {
                if ((L & 1) == 1) {
                    apply(L++, val);
                }
                if ((R & 1) == 0) {
                    apply(R--, val);
                }
            }
            pull(L0);
            pull(R0);
        }

        T query(int L, int R) {
            if (L > R) {
                return T{};
            }
            L += base_;
            R += base_;
            push(L);
            push(R);
            T left{}, right{};
            for (; L <= R; L >>= 1, R >>= 1) {
                if ((L & 1) == 1) {
                    left = query_fn_(left, tree_[L++]);
                }
                if ((R & 1) == 0) {
                    right = query_fn_(tree_[R--], right);
                }
            }
            return query_fn_(left, right);
        }

    private:
        void apply(int x, const T val) {
            tree_[x] = update_fn_(tree_[x], (static_cast<int64_t>(val) * count_[x]) % MOD);  // modified
            if (x < base_) {
                lazy_[x] = update_fn_(lazy_[x], val);
            }
        }

        void pull(int x) {
            while (x > 1) {
                x >>= 1;
                tree_[x] = query_fn_(tree_[x << 1], tree_[(x << 1) + 1]);
                if (lazy_[x]) {
                    tree_[x] = update_fn_(tree_[x], (static_cast<int64_t>(lazy_[x]) * count_[x]) % MOD);  // modified
                }
            }
        }

        void push(int x) {
            for (int h = __lg(x) - 1; h > 0; --h) {
                int y = x >> h;
                if (lazy_[y]) {
                    apply(y << 1, lazy_[y]);
                    apply((y << 1) + 1, lazy_[y]);
                    lazy_[y] = 0;
                }
            }
        }

        int base_;
        vector<T> tree_;
        vector<T> lazy_;
        vector<T> count_;  // added
        const function<T(const T&, const T&)> query_fn_;
        const function<T(const T&, const T&)> update_fn_;
    };
};

Beginner Explanation

What is Subarrays Distinct Element Sum of Squares II?

Subarrays Distinct Element Sum of Squares II (LeetCode #2916) is a Easy problem that primarily trains dynamic programming.

How to think about it

Restate the goal in your own words before coding.
Work a tiny example by hand so the invariant becomes obvious.
Identify the pattern — this problem aligns with dynamic programming, segment tree, fenwick tree, and sorted list.
Only then translate the idea into code.

Why this problem matters

It builds core muscle memory you will reuse on harder variants. Official solution notes mention: DP, Segment Tree, BIT, Fenwick Tree.

AlgoForge explanations are original teaching notes. Always open the official problem statement on LeetCode for constraints and examples.

Interview Walkthrough

Interview approach for Subarrays Distinct Element Sum of Squares II

Opening (30–60 seconds)

Clarify inputs/outputs and edge cases (empty input, single element, duplicates, overflow).
State a brute force so the interviewer knows you can solve it naively.
Propose the optimal direction tied to dynamic programming, segment tree, fenwick tree, and sorted list.

Core solution narrative

Define the state you track (pointers, DP cell, set membership, stack top, etc.).
Explain the transition when you process the next element.
Call out time (O(nlogn)) and space (O(n)) before coding.
Code cleanly; narrate variable names.

What interviewers listen for

Correctness on edge cases
Complexity honesty
Ability to discuss trade-offs (e.g., hash map space vs. sort + two pointers)

Follow-up questions they may ask

Can you solve it with less memory?
What if the input stream is infinite / doesn't fit in RAM?
How would tests look for adversarial inputs?

Optimized Approach

Optimized solution notes

The reference solutions on AlgoForge target O(nlogn) time and O(n) space.

Pattern focus: dynamic programming, segment tree, fenwick tree, and sorted list

Use the pattern as a checklist:

dynamic programming — confirm the invariant holds after each step
segment tree — confirm the invariant holds after each step
fenwick tree — confirm the invariant holds after each step
sorted list — confirm the invariant holds after each step

Multiple methods appear in the source solutions — compare them and explain when each is preferable.

Implementation tips

Prefer readable names over micro-optimizations in interviews.
Extract helpers only when they clarify (e.g., expand-around-center, DFS visit).
After AC-level logic, re-scan for off-by-one and null checks.

Complexity Analysis

Complexity

Measure	Bound
Time	O(nlogn)
Space	O(n)

How to justify this in an interview

Time: count loops, map/set operations, and recursive branching; state average vs worst case if relevant.
Space: include hash maps, recursion stack, and output allocation when the problem asks for it.

If your implementation differs from the reference, re-derive big-O from your code — never memorize a complexity you cannot defend.

Common Mistakes

Common mistakes on Subarrays Distinct Element Sum of Squares II

Skipping edge cases — empty collections, single-element inputs, max constraints.
Wrong invariant for dynamic programming, segment tree, fenwick tree, and sorted list — updating state too early or too late.
Mutating input unexpectedly when the problem forbids it.
Off-by-one in windows, ranges, or binary search bounds.
Ignoring overflow / precision for integer arithmetic problems.
Overengineering — jumping to an advanced structure when a simpler approach works.

Alternative Approaches

Alternatives

The source file includes more than one method. Compare:

Primary optimized path — best complexity for typical interviews.
Secondary approach — often brute force, sorting-based, or space-optimized variant.

Practice articulating when you would pick each (constraints, readability, follow-ups).

Edge Cases

Edge cases checklist

Minimum input size
Maximum input size / time limits
Duplicates and already-sorted input
Negative numbers / zeros (if applicable)
Disconnected structures (graphs/trees)
Single path vs branching recursion depth

Pattern Recognition

Spotting this pattern

Signal phrases that point to dynamic programming, segment tree, fenwick tree, and sorted list:

Sorted input or ability to sort without changing the answer class
Need for contiguous subarray / substring → consider sliding window
Need for O(1) membership → hash set/map
Optimal substructure + overlapping subproblems → DP
Connectivity / components → graph DFS/BFS or Union-Find

Primary topics: dynamic programming.

Follow-up Interview Questions

Follow-ups

How does the solution change if the input is a stream?
Can you solve it in-place?
What if duplicates must be handled differently?
How would you parallelize the approach?
Design tests that would break a buggy implementation.

Practice Recommendations

What to practice next

Re-solve Subarrays Distinct Element Sum of Squares II in a second language (cpp, python).
Drill 3–5 more problems tagged dynamic programming.
Teach the solution out loud in under 5 minutes.
Add this problem to your revision calendar in 3 days and 14 days.

Visualization

Conceptual diagram for Subarrays Distinct Element Sum of Squares II: show input structure (dynamic programming), highlight the moving parts of the dynamic programming, segment tree, fenwick tree, and sorted list approach, and annotate each step with the maintained invariant and complexity.

Study checklist

Read the official problem statement on LeetCode
Solve on paper / whiteboard first
Implement the dynamic programming, segment tree, fenwick tree, and sorted list approach
Verify edge cases from the checklist
State time and space complexity aloud
Compare with the AlgoForge reference solution
Schedule a revision session

Revision notes

Subarrays Distinct Element Sum of Squares II (#2916) — Easy. Pattern: dynamic programming, segment tree, fenwick tree, and sorted list. Complexity: O(nlogn) time / O(n) space. Re-derive the invariant before coding.

FAQs

What is the time complexity of Subarrays Distinct Element Sum of Squares II?+

The reference solutions aim for O(nlogn) time and O(n) space. Always re-derive complexity from the code you write in the interview.

What pattern does Subarrays Distinct Element Sum of Squares II use?+

It primarily maps to dynamic programming, segment tree, fenwick tree, and sorted list, within the broader topic of dynamic programming.

Is Subarrays Distinct Element Sum of Squares II good for interviews?+

Yes — as a Easy problem it is a solid practice target. Pair it with related problems in the same pattern family for spaced repetition.

Where can I read the official statement?+

Open the official LeetCode page for constraints and examples: https://leetcode.com/problems/subarrays-distinct-element-sum-of-squares-ii/

#	Title	Difficulty	Topics
2913	Subarrays Distinct Element Sum of Squares I	Easy	dynamic programming
2944	Minimum Number of Coins for Fruits	Medium	dynamic programming
2969	Minimum Number of Coins for Fruits II	Hard	dynamic programming
3525	Find X Value of Array II	Hard	dynamic programming
1626	Best Team With No Conflicts	Medium	dynamic programming
53	Maximum Subarray	Easy	dynamic programming
70	Climbing Stairs	Easy	dynamic programming
198	House Robber	Easy	dynamic programming