官方题解｜彩球排序

题目解析

若考虑所有球的颜色都互不相同，那么本题就变为了一道经典的「逆序对」问题。

那么我们可以考虑先求出整个序列 $X$ 的「逆序对」的数量，在考虑如何消除相同元素的球进行交换对答案产生的影响即可。

对于颜色同为 $C_i$ 的 $M_i$ 个球对应的整数序列为 $X_{p_1}, X_{p_2}, \cdots, X_{p_M}$。
令其排序以后的序列为 $X_{p_1}', X_{p_2}', \cdots, X_{p_M}'$。
那么求出上述序列的「逆序对」数量，并将其在最终的答案中减去即可。

在实现上，我们只需要创建一个「树状数组」，然后将球上的数字按照颜色分组。

对于每组内的数字，球的颜色都是相同的，所以我们对于这个序列求逆序对，在总答案中减去即可；
在结束后，可以再遍历一遍组内的数字，消除求逆序对的过程中对树状数组产生的影响。

然后我们再对整个数组求逆序对，便可得到最终的答案。

以上两个步骤的时间复杂度皆为 $\mathrm{O}(N\log{N})$。

对于 $\tt{Python}$ 而言一个可以优化的点是，倒序处理，这样每次查找就变成了“小于当前 x 的数量”，这样树状数组可以少做一半的 ask() 运算。

AC代码

C++ 代码：

#include <bits/stdc++.h>

using i64 = long long;

template<class T>
class FenwickTree {
// fenwickTree for interval [1, n]
private:
    int n;
    std::vector<T> sum;
    int lowbit(int x) {return x & -x;}
public:
    FenwickTree(int n = 0) : n(n), sum(n + 1) {}
    
    void add(int x, T d) {
        for (; x <= n; x += lowbit(x))
            sum[x] += d;
    }
    // return sum of [1, x]
    T ask(int x) {
        T res = 0;
        for (; x; x -= lowbit(x))
            res += sum[x];
        return res;
    }
    // return sum of [l, r]
    T ask(int l, int r) {
        return ask(r) - ask(l - 1);
    }
};

int main() {
    std::cin.tie(nullptr)->sync_with_stdio(false);
    int n; std::cin >> n;
    std::vector<int> c(n), x(n);
    for (auto &v : c) std::cin >> v;
    for (auto &v : x) std::cin >> v;
    std::vector<std::vector<int>> color(n + 1);
    for (int i = 0; i < n; ++i)
        color[c[i]].push_back(x[i]);
    i64 res = 0;
    FenwickTree<int> ft(n + 1);
    for (int i = 1; i <= n; ++i) {
        for (auto &val : color[i]) {
            res -= ft.ask(val + 1, n);
            ft.add(val, +1);
        }
        for (auto &val : color[i])
            ft.add(val, -1);
    }
    for (int i = 0; i < n; ++i) {
        res += ft.ask(x[i] + 1, n);
        ft.add(x[i], +1);
    }
    std::cout << res << '\n';
    return 0;
}

Python 代码：

from collections import defaultdict

class FenwickTree:
    def __init__(self, n:int = 0)->None:
        self.n = n
        self.sum = [0] * (n + 1)
    
    def add(self, p:int, x)->None:
        while p <= self.n:
            self.sum[p] += x
            p += p & -p
        
    def ask(self, p:int):
        res = 0
        while p > 0:
            res += self.sum[p]
            p -= p & -p
        return res

n = int(input())
c = list(map(int, input().split()))
x = list(map(int, input().split()))
color = defaultdict(list)
for ci, xi in zip(c, x):
    color[ci].append(xi)
res = 0
ft = FenwickTree(n + 1)
for seq in color.values():
    for xi in seq[::-1]:
        res -= ft.ask(xi - 1)
        ft.add(xi, +1)
    for xi in seq[::-1]:
        ft.add(xi, -1)
for xi in x[::-1]:
    res += ft.ask(xi - 1)
    ft.add(xi, +1)
print(res)