abc351F Double Sum

chenfy27的刷题记录 / 2024-10-09 / 原文

给定数组A[N],对所有1<=i<j<=N,计算max(A[j]-A[i],0)之和。
2<=N<=4E5; 0<=A[i]<=1E8

分析:从左到后依次处理,用平衡树维护左侧A[i],对于A[j],只需要统计值小于A[j]的那些A[i]即可,可以合并求和过程转化为前缀和。

#include <bits/stdc++.h>
using i64 = long long;

template <typename TYPE>
struct SumTreap {
    struct Node {
        TYPE data, sum;
        int rnd, siz, dup, son[2];
        Node() { data = sum = rnd = siz = dup = son[0] = son[1] = 0; }
        Node(const TYPE &d, int cnt=1) { init(d, cnt); }
        void init(const TYPE & d, int cnt) {
            data = d; dup = siz = cnt; sum = d * cnt; rnd = rand(); son[0] = son[1] = 0;
        }
    };
    SumTreap(int multi=1):multiple(multi) { reset(); }
    void setmulti(int multi) { multiple = multi; }
    int newnode(const TYPE & d, int cnt) {
        if (!reuse.empty()) {
            int r = reuse.front();
            reuse.pop_front();
            node[r].init(d, cnt);
            return r;
        }
        node.push_back({d, cnt});
        return node.size() - 1;
    }
    void reset() { root = 0; node.resize(1); reuse.clear(); }
    void maintain(int x) {
        node[x].siz = node[x].dup;
        node[x].sum = node[x].data * node[x].dup;
        if (node[x].son[0]) {
            node[x].siz += node[node[x].son[0]].siz;
            node[x].sum += node[node[x].son[0]].sum;
        }
        if (node[x].son[1]) {
            node[x].siz += node[node[x].son[1]].siz;
            node[x].sum += node[node[x].son[1]].sum;
        }
    }
    void rotate(int d, int &r) {
        int k = node[r].son[d^1];
        node[r].son[d^1] = node[k].son[d];
        node[k].son[d] = r;
        maintain(r);
        maintain(k);
        r = k;
    }
    void insert(const TYPE &data, int cnt, int &r) {
        if (r) {
            if (!(data < node[r].data) && !(node[r].data < data)) {
                if (multiple) {
                    node[r].dup += cnt;
                    maintain(r);
                }
            } else {
                int d = data < node[r].data ? 0 : 1;
                int u = node[r].son[d];
                insert(data, cnt, u);
                node[r].son[d] = u;
                if (node[node[r].son[d]].rnd > node[r].rnd) {
                    rotate(d^1, r);
                } else {
                    maintain(r);
                }
            }
        } else {
            r = newnode(data, cnt);
        }
    }
    TYPE kth(int k) {
        assert(1 <= k && k <= size());
        for (int r = root; r; ) {
            int x = node[r].son[0] ? node[node[r].son[0]].siz : 0;
            int y = node[r].dup;
            if (k <= x) {
                r = node[r].son[0];
            } else if (k <= x + y) {
                return node[r].data;
            } else {
                k -= x + y;
                r = node[r].son[1];
            }
        }
        assert(k == 0);
        return 0;
    }
    TYPE Kth(int k) {
        assert(1 <= k && k <= size());
        for (int r = root; r; ) {
            int x = node[r].son[1] ? node[node[r].son[1]].siz : 0;
            int y = node[r].dup;
            if (k <= x) {
                r = node[r].son[1];
            } else if (k <= x + y) {
                return node[r].data;
            } else {
                k -= x + y;
                r = node[r].son[0];
            }
        }
        assert(k == 0);
        return 0;
    }
    TYPE ksum(int k) {
        assert(0 <= k && k <= node[root].siz);
        TYPE ans = 0;
        for (int r = root; r && k; ) {
            int x = node[r].son[0] ? node[node[r].son[0]].siz : 0;
            int y = node[r].dup;
            if (k <= x) {
                r = node[r].son[0];
            } else if (k <= x + y) {
                ans += node[node[r].son[0]].sum + node[r].data * (k - x);
                k = 0;
            } else {
                ans += node[node[r].son[0]].sum + node[r].data * y;
                k -= x + y;
                r = node[r].son[1];
            }
        }
        return ans;
    }
    TYPE kSum(int k) {
        assert(0 <= k && k <= node[root].siz);
        TYPE ans = 0;
        for (int r = root; r && k; ) {
            int x = node[r].son[1] ? node[node[r].son[1]].siz : 0;
            int y = node[r].dup;
            if (k <= x) {
                r = node[r].son[1];
            } else if (k <= x + y) {
                ans += node[node[r].son[1]].sum + node[r].data * (k - x);
                k = 0;
            } else {
                ans += node[node[r].son[1]].sum + node[r].data * y;
                k -= x + y;
                r = node[r].son[0];
            }
        }
        return ans;
    }
    void erase(const TYPE& data, int cnt, int & r) {
        if (r == 0) return;
        int d = -1;
        if (data < node[r].data) {
            d = 0;
        } else if (node[r].data < data) {
            d = 1;
        }
        if (d == -1) {
            node[r].dup -= cnt;
            if (node[r].dup > 0) {
                maintain(r);
            } else {
                if (node[r].son[0] == 0) {
                    reuse.push_back(r);
                    r = node[r].son[1];
                } else if (node[r].son[1] == 0) {
                    reuse.push_back(r);
                    r = node[r].son[0];
                } else {
                    int dd = node[node[r].son[0]].rnd > node[node[r].son[1]].rnd ? 1 : 0;
                    rotate(dd, r);
                    erase(data, cnt, node[r].son[dd]);
                }
            }
        } else {
            erase(data, cnt, node[r].son[d]);
        }
        if (r) maintain(r);
    }
    int ltcnt(const TYPE &data) {
        int ans = 0;
        for (int r = root; r; ) {
            int x = node[r].son[0] ? node[node[r].son[0]].siz : 0;
            if (node[r].data < data) {
                ans += node[r].dup + x;
                r = node[r].son[1];
            } else if (data < node[r].data) {
                r = node[r].son[0];
            } else {
                ans += x;
                break;
            }
        }
        return ans;
    }
    int gtcnt(const TYPE &data) {
        int ans = 0;
        for (int r = root; r; ) {
            int x = node[r].son[1] ? node[node[r].son[1]].siz : 0;
            if (node[r].data < data) {
                r = node[r].son[1];
            } else if (data < node[r].data) {
                ans += node[r].dup + x;
                r = node[r].son[0];
            } else {
                ans += x;
                break;
            }
        }
        return ans;
    }
    int count(const TYPE &data) {
        for (int r = root; r; ) {
            if (data < node[r].data) {
                r = node[r].son[0];
            } else if (node[r].data < data) {
                r = node[r].son[1];
            } else {
                return node[r].dup;
            }
        }
        return 0;
    }
    std::pair<bool,TYPE> prev(const TYPE &data) {
        std::pair<bool,TYPE> ans = {false, 0};
        for (int r = root; r; ) {
            if (node[r].data < data) {
                if (ans.first) {
                    ans.second = std::max(ans.second, node[r].data);
                } else {
                    ans = {true, node[r].data};
                }
                r = node[r].son[1];
            } else {
                r = node[r].son[0];
            }
        }
        return ans;
    }       
    std::pair<bool,TYPE> next(const TYPE &data) {
        std::pair<bool,TYPE> ans = {false, 0};
        for (int r = root; r; ) {
            if (data < node[r].data) {
                if (ans.first) {
                    ans.second = std::min(ans.second, node[r].data);
                } else {
                    ans = {true, node[r].data};
                }
                r = node[r].son[0];
            } else {
                r = node[r].son[1];
            }
        }
        return ans;
    }
    std::vector<Node> node;
    std::deque<int> reuse;
    int root, multiple;
    void insert(const TYPE& data, int cnt=1) { insert(data, cnt, root); }
    void erase(const TYPE& data, int cnt=1) { erase(data, cnt, root); }
    int size() const { return root ? node[root].siz : 0; }
    int lecnt(const TYPE& data) { return size() - gtcnt(data, root); }
    int gecnt(const TYPE& data) { return size() - ltcnt(data, root); }
};

void solve() {
    int N;
    std::cin >> N;
    SumTreap<i64> tr;
    std::vector<int> A(N);
    for (int i = 0; i < N; i++) {
        std::cin >> A[i];
    }

    i64 ans = 0;
    for (int i = 0; i < N; i++) {
        i64 cnt1 = tr.ltcnt(A[i]);
        ans += cnt1 * A[i] - tr.ksum(cnt1);
        tr.insert(A[i]);
    }
    std::cout << ans << "\n";
}

int main() {
    std::cin.tie(0)->sync_with_stdio(0);
    std::cout << std::fixed << std::setprecision(10);
    int t = 1;
    // std::cin >> t;
    while (t--) solve();
    return 0;
}