kmjp's blog

競技プログラミング参加記です

yukicoder : No.2115 Making Forest Easy

Easyは割と典型か。
https://yukicoder.me/problems/no/2115

問題

木を成すN頂点の無向グラフが与えられる。
各点vには整数A[v]が書かれている。

このグラフの辺をk本切断すると、(k+1)個の木ができる。
このグラフのスコアは、各連結成分における(頂点に書かれた整数値の最大値)×(頂点数)の総和を取ったものとする。
辺の切断のしかた2^(N-1)通りに対し、上記スコアの総和を求めよ。

解法

まず頂点番号を値の降順に並べ替えておく。
i番の頂点に対し、「i番以降の頂点だけからなる連結成分の作り方と、その成分の頂点数の積の総和」をS(i)とすると、S(i)*A[i]の総和を取れば解になる。
iごとにO(N)かけてS(i)を求めよう。

i番の頂点を根とする木を考え、木DPを行おう。
各点vの状態としては、

  • vを含む連結成分に、i番以前の頂点を含まない場合における
    • 辺の切断方法の組み合わせ
    • (辺の切断方法の組み合わせ)*(vを含む連結成分のサイズ)の総和
  • vを含む連結成分に、i番以前の頂点を含む場合における
    • 以下同上

を持てばよい。

int N;
int A[101010];
int D[101010];
int C[101010];
vector<int> E[101010];
const ll mo=998244353;

ll dp1[101010][2];
ll dp2[101010][2];


ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

void dfs(int cur,int pre,int sc) {
	ll a[2]={},b[2]={};
	a[C[cur]>sc]=b[C[cur]>sc]=1;
	int i,j;
	FORR(e,E[cur]) if(e!=pre) {
		dfs(e,cur,sc);
		ll na[2]={},nb[2]={};
		// not connect
		FOR(i,2) {
			na[i]+=a[i]*(dp1[e][0]+dp1[e][1]);
			nb[i]+=b[i]*(dp1[e][0]+dp1[e][1]);
		}
		// connect
		FOR(i,2) FOR(j,2) {
			na[max(i,j)]+=a[i]*dp1[e][j];
			nb[max(i,j)]+=b[i]*dp1[e][j]+a[i]*dp2[e][j];
		}
		
		FOR(i,2) {
			a[i]=na[i]%mo;
			b[i]=nb[i]%mo;
		}
	}
	FOR(i,2) {
		dp1[cur][i]=a[i];
		dp2[cur][i]=b[i];
	}
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	vector<pair<int,int>> V;
	FOR(i,N) {
		cin>>A[i];
		V.push_back({A[i],i});
	}
	sort(ALL(V));
	FOR(i,V.size()) C[V[i].second]=i;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	
	ll ret=0;
	FOR(i,N) {
		dfs(i,i,C[i]);
		ret+=A[i]*dp2[i][0]%mo;
	}
	cout<<ret%mo<<endl;
	
}

まとめ

今年は調子よく行けるといいんだけどね。