subarrays-distinct-element-sum-of-squares-ii.cpp

// 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_;
    };
};