891C - Envy - CodeForces Solution

#include<iostream>
#include<cstring>
#include<algorithm>
#include<vector>
#include<map>
#include<cmath>
#include<set>
#define PII pair<int, int>
#define LL int 
#define PLL pair<LL, LL>
using namespace std;
#define N 500010
LL f[N],si[N];

vector<pair<LL&,LL>>his_fa,his_si;
map<LL,vector<LL> >pos[N];
vector<LL>g[N];
LL ans[N];
struct node
{
    LL u,v,w;
}a[N];
LL n,m;
LL find(LL x)
{
    while(x!=f[x])x=f[x];
    return f[x];
}
void merge(LL x,LL y)
{
    x=find(x);
    y=find(y);
    if(x!=y)
    {
        if(si[x]<si[y])swap(x,y);
        his_fa.push_back({f[y],f[y]});
        f[y]=x;
        his_si.push_back({si[y],si[y]});
        si[x]+=si[y];
    }
}
LL history()
{
    return his_fa.size();
}
void roll(LL h)
{
    while(his_fa.size()>h)
    {
        his_fa.back().first=his_fa.back().second;
        his_fa.pop_back();
        his_si.back().first=his_si.back().second;
        his_si.pop_back();
    }
}

void solve()
{
    scanf("%d%d",&n,&m);
    for(int i=1;i<=n;i++)
    {
        f[i]=i;
        si[i]=1;
    }
    for(int i=1;i<=m;i++)
    {
        scanf("%d%d%d",&a[i].u,&a[i].v,&a[i].w);
        g[a[i].w].push_back(i);
    }
    LL q;
    scanf("%d",&q);
    for(int i=1;i<=q;i++)
    {
        LL k;
        scanf("%d",&k);
        for(int j=1;j<=k;j++)
        {
            LL idx;
            scanf("%d",&idx);
            pos[a[idx].w][i].push_back(idx);
        }
    }
    for(int i=1;i<=5e5;i++)
    {
        for(auto it:pos[i])
        {
            LL v=it.first;
            LL h=history();
            for(auto j:it.second)
            {
                if(find(a[j].u)==find(a[j].v))ans[v]=-1;
                merge(a[j].u,a[j].v);
            }
            roll(h);
        }
        for(auto it:g[i])merge(a[it].u,a[it].v);
    }
    for(int i=1;i<=q;i++)
        if(ans[i]==0)printf("YES\n");
        else printf("NO\n");
}
int main()
{
    solve();
    return 0;
}