AtCoder ABC #235 : Ex - Painting Weighted Graph

方針は悪くなかったけど、計算量の見積もりが甘かったな…。
https://atcoder.jp/contests/abc235/tasks/abc235_h

問題

無向グラフが与えられる。
各辺には重みが設定されている。

初期状態では各点は黒く塗られている。
ここで、以下の処理を行う。

• 頂点vと重みxを指定する。vから重みx以下で到達できる全頂点を赤く塗る

上記処理をK回以下行うとき、頂点の塗り分け方は何通りか。

解法

ある連結成分Sについて、多項式P(S,x)を、xのn次の係数は、S内のうち最低n回処理を行わないと実現できない塗り分け方の組み合わせとする。

あとはクラスカル法の要領で、連結成分をさらに連結していく。
もしクラスカル法の過程で、ある等しい重みの辺によりm個の連結成分s0,s1,s2,...が1つの連結成分Sに連結されるとする。
この時、
P(S,x) = Prod(P(s_i,x))-x^m+x
と表現できる。
(-x^m+x)の部分の意味だが、もし元の連結成分が全部赤い場合、それまでの個々の連結成分内で1回ずつ処理を行うように数え上げていたが、そもそもそれらを連結後に1回処理を行えば同じ塗り分けが再現できるので、m回塗る組み合わせを1回引き、1回塗る組み合わせを1回加算している。

int N,M,K;
int A[505050],B[505050],C[505050];

template<int um> class UF {
	public:
	vector<int> par,rank,cnt;
	UF() {par=rank=vector<int>(um,0); cnt=vector<int>(um,1); for(int i=0;i<um;i++) par[i]=i;}
	void reinit(int num=um) {int i; FOR(i,num) rank[i]=0,cnt[i]=1,par[i]=i;}
	int operator[](int x) {return (par[x]==x)?(x):(par[x] = operator[](par[x]));}
	int count(int x) { return cnt[operator[](x)];}
	int operator()(int x,int y) {
		if((x=operator[](x))==(y=operator[](y))) return x;
		cnt[y]=cnt[x]=cnt[x]+cnt[y];
		if(rank[x]>rank[y]) return par[x]=y;
		rank[x]+=rank[x]==rank[y]; return par[y]=x;
	}
};
UF<505050> uf;
vector<ll> dp[101010];
const ll mo=998244353;
int num[505050];

vector<ll> MultPoly(vector<ll> P,vector<ll> Q) {
	int a=P.size()-1;
	int b=Q.size()-1;
	int c=min(K,a+b);
	vector<ll> R(c+1);
	
	for(a=0;a<P.size();a++) {
		for(b=0;b<Q.size()&&a+b<R.size();b++) {
			(R[a+b]+=P[a]*Q[b])%=mo;
		}
	}
	
	return R;
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>K;
	map<int,vector<int>> Es;
	FOR(i,M) {
		cin>>A[i]>>B[i]>>C[i];
		A[i]--,B[i]--;
		Es[C[i]].push_back(i);
	}
	FOR(i,N) dp[i]={1,1};
	
	FORR2(c,E,Es) {
		set<int> Ts;
		FORR(i,E) {
			x=uf[A[i]];
			y=uf[B[i]];
			if(x!=y) {
				num[x]=num[y]=1;
				Ts.insert(x);
				Ts.insert(y);
				uf(x,y);
			}
		}
		FORR(e,Ts) if(uf[e]!=e) {
			x=uf[e];
			num[x]++;
			dp[x]=MultPoly(dp[x],dp[e]);
		}
		FORR(e,Ts) if(uf[e]==e) {
			if(num[e]<dp[e].size()) (dp[e][num[e]]+=mo-1)%=mo;
			dp[e][1]++;
		}
	}

	vector<ll> ret={1};
	FOR(i,N) if(uf[i]==i) ret=MultPoly(ret,dp[i]);
	ll pat=0;
	FORR(k,ret) pat+=k;
	cout<<pat%mo<<endl;
	
	
}

まとめ

本番無駄にNTTとかやってTLEしてしまった。