algoNote

プログラミング関連

Codeforces Round #448 (Div. 2) C - Square Subsets

コンテスト中にAC出来ませんでしたが、解けたので好きな問題です(オイ)

問題

codeforces.com

N 個の要素から成る数列 {a_i} が与えられます。
要素の積が k^{2} の形で表せるような部分集合の取り出し方は何通り存在しますか。109 + 7で割った余りを求めてください。

制約

  • {1 \leq N \leq 10^{5}}
  • {1 \leq a_i \leq 70}

考察

まず、積が2乗の形になるというのは、積を素因数分解したときの指数部分がすべて偶数になるということです。
よって条件を満たす部分集合は、要素をそれぞれ素因数分解したとき、各素因数の指数部分の和がすべて偶数になるようなものです。

ここで、交わらない2つの部分集合の和集合がどうなるかを考えてみます。
集合 A の要素の積を {P_A} として、集合 B の要素の積を {P_B} とします。

今、

{P_A = p_1^{s} \cdot p_2^{t} \cdot p_3^{u} \cdot \cdots}

{P_B = p_1^{x} \cdot p_2^{y} \cdot p_3^{z} \cdot \cdots}

と表せたとします。 ({p_i} はi番目の素数です。)

AB の和集合の要素の積 {P_{A \bigcup B}}

{P_{A \bigcup B} = p_1^{s + x} \cdot p_2^{t + y} \cdot p_3^{u + z} \cdot \cdots}

となります。

P_{A \bigcup B} が条件を満たす部分集合になるとき、 {s+x}, {t+y}, {u+z}, ... は偶数です。
和が偶数であればいいので、sが奇数ならxは奇数ですし、sが偶数ならxは偶数です。

これから分かるように、素因数分解したときの指数部分の偶奇だけが重要です。 70以下の素因の個数を調べてみると高々19個しか存在しないことが分かるので、これをビットで持つことが出来ます。
ここでは、i 番目のビットが0なら {p_i}の指数部分が偶数で1なら奇数であるとします。
ビットで持たせることで、その偶奇の組になる部分集合が何通りあるかを持ちながらbitDPの要領で答えを計算できそうです。

以下 {bit(S)} は集合 S に対応するビット列を表すことにします。

具体的なDPの遷移は、
ある集合 S に値 x を加えるなら dp[次][bit(S) ^ bit(x)] += dp[今][bit(S)] とできそうです。
(1の部分の素因数の個数の偶奇が入れ替わるのでXORをとります。)

しかし、配列のすべての要素について上の遷移をしようとすると、105 * 220 > 1011 なのでTLEしてしまいます。

そこでもう少し考えると、同じ値についてはまとめて考えられることが分かります。

ある集合 S に値 x を複数加える場合

  • 偶数個加える場合 dp[次][bit(S)] += dp[今][bit(S)]
  • 奇数個加える場合 dp[次][bit(S) ^ bit(x)] += dp[今][bit(S)]

となるので、

  • K(K>0) から偶数個取る組み合わせの数 { \displaystyle \sum_{i=0}^{\lfloor \frac{K}{2} \rfloor} {}_{K}C_{2i} = 2^{K-1} }
  • K(K>0) から奇数個取る組み合わせの数 { \displaystyle \sum_{i=0}^{\lfloor \frac{K-1}{2} \rfloor} {}_{K}C_{2i+1} = 2^{K-1} }

となることを用いて、K 個の x について

  • 偶数個 dp[次][bit(S)] += dp[今][bit(S)] * 2^{K-1}
  • 奇数個 dp[次][bit(S) ^ bit(x)] += dp[今][bit(S)] * 2^{K-1}

と、まとめて処理出来ます。

1から70までの個数を事前にカウントしておくことで、 70 * 220 < 108 でDPが出来ます。

ソースコード

2のcnt[i]-1乗する部分は二分累乗法でやってますが、線形でやっても通ると思います。

空集合は含めないので最後に1引いています。