FPS24 I - スコア 勉強メモ

問題リンク

https://atcoder.jp/contests/fps-24/tasks/fps_24_i

問題文

整数からなる要素数 $N$ の集合 $A = \{A_1, A_2, \cdots, A_N\}$ が与えられます。

$A$ の要素数 $K$ の部分集合 $B = \{B_1, B_2, \cdots, B_K\}$ にすべてに対する $\prod_{1 \leq i \leq K} B_i$ （スコア）の総和を $998244353$ で割った余りを求めてください。

制約

• 入力される値はすべて整数
• $1 \leq N \leq 2 \times 10^5$
• $1 \leq K \leq N$
• $1 \leq A_i \leq 10^8$
• $j \ne j \Longrightarrow A_i \ne A_j$

解説

$[x^k]f$ を $A$ から要素を $k$ 個選んだときのスコアの総和としたとき、

$$f =\prod_{i = 1}^N (1 + A_i x)$$

となります。

$x$ を掛けることを、$1$ 要素を選択していることが対応していると解釈すると、見通しが良いかもしれません。そして、所望の答えは $[x^K]f$ です。

これは畳み込みの順序を マージテク することで、計算量 $O(N \log^2 N)$ で計算することができます。

このマージテクは、以下のqueueを使ったアルゴリズムで簡単に実装できます。

• queue を $(f_1, f_2, \cdots, f_N)$ と初期化する。 $(f_i := 1 + A_i x)$
• queue の長さが $1$ になるまで以下を繰り返す。
  • queue の状態を $(f_1, f_2, f_3, \cdots, f_m) \to (f_3, f_4, \cdots, f_m, f_1 * f_2)$ に更新する。(直前の queue の長さを $m$ とする。)
• queue の唯一の要素 $(f_1 * f_2 * \cdots * f_N)$ が求めたい多項式である。

各操作について次数が最小である $2$ つの多項式について畳み込めているので、マージ回数 $O(\log N)$ 回で済んでいることが分かります。

実装例（Python3）

# multiply(f, g) : 多項式 f と g を畳み込む

N, K = [int(s) for s in input().split()]
A = [int(s) for s in input().split()]

que = deque()
for a in A:
    que.append([1, a])

while len(que) > 1:
    f, g = que.popleft(), que.popleft()
    res = multiply(f, g)
    que.append(res)

res = que[0]
print(res[K])

https://atcoder.jp/contests/fps-24/submissions/75298964

所感

こんなところにもマージテクが出てくるのは感動した。天才っているもんだ。

$f$ の立式が思いつかずクネクネしていたので、立式の部分も強化しないとなぁ。