AtCoder ABC #364 (日本レジストリサービス（JPRS）プログラミングコンテスト2024#2) : G - Last Major City

考察が甘くてTLE連発してしまった。
https://atcoder.jp/contests/abc364/tasks/abc364_g

問題

連結無向グラフが与えられる。
各辺には重みが与えられる。

各点vについて、点1～(K-1)及びvを連結とするために必要な辺の重みの総和の最小値を求めよ。

解法

dp(v,mask) := 点vと、点1～点(K-1)のうちmaskが示す部分集合が連結となるための辺の重みの総和の最小値
とする。解は、maskが点1～(K-1)をすべて含む時のdp(v,mask)となる。

dp(v,mask1|mask2) = min(dp(v,mask1)+dp(u,mask2)+dist(u,v))
としてBitDPすると解けるが、これだとO(3^K*N^2)かかりTLEする。

vに対しuをvの隣接点だけたどるようにするのと、mask毎にdp(v,mask) = min(dp(v,mask),dp(u,mask)+dist(u,v))と毎回ダイクストラ法を動かすと、O(3^K*NlogN)になって間に合う。

int N,M,K;
vector<pair<int,int>> E[4040];

ll dp[4040][4040];
ll dp2[4040][1<<9];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>K;
	K--;
	FOR(i,M) {
		cin>>x>>y>>k;
		E[x-1].push_back({y-1,k});
		E[y-1].push_back({x-1,k});
	}
	FOR(i,K) {
		FOR(j,N) dp[i][j]=1LL<<60;
		dp[i][i]=0;
		priority_queue<pair<ll,int>> Q;
		Q.push({0,i});
		while(Q.size()) {
			ll co=-Q.top().first;
			int cur=Q.top().second;
			Q.pop();
			if(dp[i][cur]!=co) continue;
			FORR2(e,c,E[cur]) if(chmin(dp[i][e],co+c)) Q.push({-dp[i][e],e});
		}
	}
	int mask,mask2;
	FOR(mask,1<<K) FOR(i,N) {
		FOR(j,K) if(mask&(1<<j)) dp2[i][mask]+=dp[j][i];
	}
	FOR(mask,1<<K) {
		for(int sm=mask;sm>=0;sm--) {
			sm&=mask;
			int sm2=mask^sm;
			FOR(i,N) FORR2(j,c,E[i]) dp2[i][mask]=min(dp2[i][mask],dp2[i][sm]+dp2[j][sm2]+c);
		}
		priority_queue<pair<ll,int>> Q;
		FOR(i,N) Q.push({-dp2[i][mask],i});
		while(Q.size()) {
			ll co=-Q.top().first;
			int cur=Q.top().second;
			Q.pop();
			if(dp2[cur][mask]!=co) continue;
			FORR2(e,c,E[cur]) if(chmin(dp2[e][mask],co+c)) Q.push({-dp2[e][mask],e});
		}
		
	}
	for(i=K;i<N;i++) cout<<dp2[i][(1<<K)-1]<<endl;
	
	
}

まとめ

計算量の見積もりが甘かったな…。