Atcoder Grand Contest 012 D. Colorful Balls

Atcoder Grand Contest 012 D. Colorful Balls

感想

前半の考察はできたのに、ほとんど同じの後半の発想がなぜかできなかった。

解法

まずすべての色が等しければ、答えは$\ 1\ $となる。以下は2種類以上の色があるものとする。

操作1では色の配置を変えることができないので、この操作の役割というのはできるだけ操作2ができるような状態にもっていくことである。操作2をするときには、常に操作1によって、その2つのボールの重さをできるだけ小さくなるようにしておくのが最適である。このように考えると結局、すべてのボールについて、そのボールと同じ色であって、かつ交換可能なボールの中で重さの最小値をとって、それをそのボールの重さとしてしまって良い。するとあとは操作2のみを考えればよい。

以下、上の操作によって変化した重さで考える。重さが最も小さいボールに注目すると、このボールと色が異なっていて、かつこのボールと交換可能ではないボールは、他のどのボールとも交換することができない。逆に、このボールと色が異なっていて、交換可能であるならば、それらのボールは、すべて互いに交換可能である。それは最小のボールを仲介にして、スワップを3回することによって、他のボールの位置を全く変えずに任意の2個をスワップできるからである。これらの並び順は重複を含む順列なので簡単に計算できる。

上の重さが最も小さいボールと同じ色のボールに対してはどう考えればよいだろうか？それはやはり上のボールとは異なる色の中で、重さが最小のものをとってくることで、同様に考えることができる。

この2つのボールがもし交換不可能であるならば、どの2つのボールも交換不可能であるので、答えは$\ 1\ $である。交換可能であるならば、これらの和集合の任意の2つは交換可能ということになる。

実装

#include "bits/stdc++.h"
using namespace std;

#define all(x)  x.begin(), x.end()
#define mp      make_pair
#define pll     pair<long long, long long>

static const long long INFL = 0x3f3f3f3f3f3f3f3fLL;
static const int INF = 0x3f3f3f3f;
static const int MOD = 1000000007;

struct UnionFind {
        int n;
        vector<int> parent;
        vector<int> rank;
        vector<int> num;
        int find(int x) {
                if (parent[x] == x) return  x;
                return parent[x] = find(parent[x]);
        }
        UnionFind(int n_) {
                n = n_;
                parent.resize(n);
                for (int i = 0; i < n; i ++) parent[i] = i;
                rank.assign(n, 0);
                num.assign(n, 1);
        }
        void unite(int x, int y) {
                if ((x = find(x)) != (y = find(y))) {
                        if (rank[x] < rank[y]) {
                                parent[x] = y;
                                num[y] += num[x];
                        } else {
                                parent[y] = x;
                                if (rank[x] == rank[y]) rank[x] ++;
                                num[x] += num[y];
                        }
                        n --;
                }
        }
        bool same(int x, int y) { return find(x) == find(y); }
        int get() { return n; }
        int get(int x) { return num[find(x)]; }
};

long long inv(long long a) {
        long long res = 1;
        long long n = MOD - 2;
        while (n > 0) {
                if (n & 1) res = res * a % MOD;
                a = a * a % MOD;
                n >>= 1;
        }
        return res;
}
long long fact[202020];

void init() {
        fact[0] = 1;
        for (int i = 1; i < 202020; i ++) fact[i] = fact[i - 1] * i % MOD;
}

void solve() {
        init();
        long long n, x, y;
        cin >> n >> x >> y;
        vector<vector<pll>> color(200000);
        for (int i = 0; i < n; i ++) {
                long long c, w;
                cin >> c >> w;
                color[c - 1].emplace_back(i, w);
        }
        int cnt = 0;
        for (int i = 0; i < n; i++) if (color[i].size() > 0) cnt ++;
        if (cnt == 1) {
                cout << 1 << endl;
                return;
        }
        vector<pll> ball(n);
        for (int i = 0; i < 200000; i ++) if (!color[i].empty()) {
                long long minwei = INFL;
                for (int j = 0; j < color[i].size(); j ++) minwei = min(minwei, color[i][j].second);
                for (int j = 0; j < color[i].size(); j ++) {
                        ball[color[i][j].first] = mp(i, (minwei + color[i][j].second <= x ? minwei : color[i][j].second));
                }
        }       
        int mi1 = INF, idx1, mi2 = INF, idx2;
        for (int i = 0; i < n; i ++) {
                if (ball[i].second < mi1) {
                        mi1 = ball[i].second;
                        idx1 = i;
                }
        }
        for (int i = 0; i < n; i ++) {
                if (ball[i].first != ball[idx1].first && ball[i].second < mi2) {
                        mi2 = ball[i].second;
                        idx2 = i;
                }
        }
        UnionFind uf(n);
        for (int i = 0; i < n; i ++) {
                if (ball[i].second + ball[idx1].second <= y && ball[i].first != ball[idx1].first) uf.unite(i, idx1);
                if (ball[i].second + ball[idx2].second <= y && ball[i].first != ball[idx2].first) uf.unite(i, idx2);
        }
        if (!uf.same(idx1, idx2)) {
                cout << 1 << endl;
        } else {
                int all = uf.get(idx1);
                map<int, int> color_cnt;
                for (int i = 0; i < n; i ++) if (uf.same(i, idx1)) color_cnt[ball[i].first] ++;
                long long ans = fact[all];
                for (auto c : color_cnt) ans = ans * inv(fact[c.second]) % MOD;
                cout << ans % MOD << endl;
        }
        return;
}

int main() {
        ios_base::sync_with_stdio(false);
        cin.tie(0);
        solve();
        return 0;
}