Codeforces 891C Envy

Envy

感觉这种最小生成树上的啥题都差不多的解法。。

#include<bits/stdc++.h>

#define LL long long

#define fi first

#define se second

#define mk make_pair

#define PLL pair<LL, LL>

#define PLI pair<LL, int>

#define PII pair<int, int>

#define SZ(x) ((int)x.size())

#define ull unsigned long long

using namespace std;

const int N = 5e5 + ;

const int inf = 0x3f3f3f3f;

const LL INF = 0x3f3f3f3f3f3f3f3f;

const int mod = 1e9 + ;

const double eps = 1e-;

const double PI = acos(-);

int n, m, q, a[N], b[N], w[N], va[N], vb[N], e[N];

vector<int> vc[N];

int fa[N];

int getRoot(int x) {

    return fa[x] == x ? x : fa[x] = getRoot(fa[x]);

}

bool cmp(const int& a, const int& b) {

    return w[a] < w[b];

}

int main() {

    scanf("%d%d", &n, &m);

    for(int i = ; i <= n; i++) fa[i] = i;

    for(int i = ; i <= m; i++) {

        scanf("%d%d%d", &a[i], &b[i], &w[i]);

        vc[w[i]].push_back(i);

    }

    for(int i = ; i <= ; i++) {

        if(!SZ(vc[i])) continue;

        for(auto& x : vc[i]) {

            va[x] = getRoot(a[x]);

            vb[x] = getRoot(b[x]);

        }

        for(auto& x : vc[i]) {

            fa[getRoot(b[x])] = getRoot(a[x]);

        }

    }

    scanf("%d", &q);

    while(q--) {

        bool flag = true;

        int num; scanf("%d", &num);

        for(int i = ; i <= num; i++) scanf("%d", &e[i]);

        for(int i = ; i <= num; i++) {

            if(va[e[i]] == vb[e[i]]) {

                flag = false;

                break;

            }

        }

        if(!flag) {

            puts("NO");

            continue;

        }

        sort(e + , e +  + num, cmp);

        for(int L = , R = ; L <= num && flag; L = R) {

            while(R <= num && w[e[R]] == w[e[L]]) R++;

            for(int i = L; i < R; i++) {

                int x = va[e[i]], y = vb[e[i]];

                fa[x] = x, fa[y] = y;

            }

            for(int i = L; i < R; i++) {

                int x = va[e[i]], y = vb[e[i]];

                int u = getRoot(x), v = getRoot(y);

                if(u == v) {

                    flag = false;

                    break;

                }

                fa[v] = u;

            }

        }

        if(flag) puts("YES");

        else puts("NO");

    }

    return ;

}

/*

*/