Atcoder Grand Contest 002 D. Stamp Rally

Atcoder Grand Contest 002 D. Stamp Rally

D: Stamp Rally - AtCoder Grand Contest 002 | AtCoder

解法

二分探索をまとめてやる感じのやつ。各$\ mid\ $の値に関するソートをしておくことで左から順に見ていける。$\ O(m {\log}^2 m)\ $。

実装

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

#define all(x)  x.begin(), x.end()
#define mp      make_pair

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)]; }
};

int main() {
        int n, m;
        scanf("%d%d", &n, &m);
        vector<pair<int, int>> es(m);
        for (int i = 0; i < m; i ++) {
                int a, b;
                scanf("%d%d", &a, &b);
                a --, b --;
                es[i] = mp(a, b);
        }
        int q;
        scanf("%d", &q);
        vector<int> x(q), y(q), z(q);
        vector<int> lb(q), ub(q);
        for (int i = 0; i < q; i ++) {
                scanf("%d%d%d", &x[i], &y[i], &z[i]);
                x[i] --, y[i] --;
                lb[i] = 0, ub[i] = m;
        }
        vector<int> idx;
        for (int i = 0; i < q; i ++) idx.push_back(i);
        for (int loop = 0; loop < 20; loop ++) {
                sort(all(idx), [&](const int &a, const int &b) {
                        return (lb[a] + ub[a]) / 2 < (lb[b] + ub[b]) / 2;
                });
                UnionFind uf(n);
                int p = 0;
                for (auto &i : idx) {
                        int mid = (lb[i] + ub[i]) / 2;
                        while (p <= mid) {
                                uf.unite(es[p].first, es[p].second);
                                p ++;
                        }
                        int sum = uf.get(x[i]) + uf.get(y[i]);
                        if (uf.same(x[i], y[i])) sum /= 2;
                        if (sum < z[i]) { 
                                lb[i] = mid + 1;
                        } else { 
                                ub[i] = mid;
                        }
                }
        }
        for (int i = 0; i < q; i ++) printf("%d\n", lb[i] + 1);
        return 0;
}