yukicoder : No.1912 Get together 2

1911よりこちらの方がすんなりだった。
https://yukicoder.me/problems/no/1912

問題

N匹の区別あるスライムをM個の区別ある箱に入れる。
各スライムiの価値と、スライムを入れられる箱の集合が与えられる。

各箱の価値は、中に入れたスライムの総価値の二乗となる。
最適なスライムの入れ方をしたとき、各箱の総価値の最大値を求めよ。

解法

まず1段階目として、以下を求める。
S(mask) := maskに該当する箱のいずれかに入れられるスライムの価値の総和
これは、
T(mask) := maskに該当する箱のいずれにも入れられないスライムの価値の総和
とするとT(mask)は高速ゼータ変換でO(N+M*2^M)で求めることができ、S(mask) = sum(価値の総和)-T(mask)で計算できる。

次に、
dp(mask) := maskに該当する箱のいずれかに入れられるスライムを最適に入れたとき、各箱の価値の総和
とする。
maskに含まれるbitの1つをvとしたとき、(mask^v)の箱には入れられないけど、vの箱には入れられるスライムの価値の総和を考えると、
dp(mask) = max(dp(mask^v) + (S(mask)-S(mask^v))^2)
というように遷移できる。

int N,M;
int V[101010];
int A[101010];

int S[1<<20];
ll dp[1<<20];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	ll sum=0;
	FOR(i,N) {
		cin>>V[i];
		sum+=V[i];
	}
	FOR(i,N) {
		cin>>s;
		FOR(j,M) if(s[j]=='o') A[i]|=1<<j;
		S[((1<<M)-1)^A[i]]+=V[i];
	}
	FOR(i,M) {
		FOR(j,1<<M) if(j&(1<<i)) S[j^(1<<i)]+=S[j];
	}
	FOR(j,1<<M) S[j]=sum-S[j];
	
	int mask;
	FOR(mask,1<<M) if(mask) {
		FOR(i,M) if(mask&(1<<i)) {
			ll a=S[mask]-S[mask^(1<<i)];
			dp[mask]=max(dp[mask],dp[mask^(1<<i)]+a*a);
		}
	}
	
	cout<<dp[(1<<M)-1]<<endl;
}

まとめ

S(mask)を直接求めようとしてちょっと手間取ってしまった。