Kazun の競プロ記録

競技プログラミングに関する様々な話題を執筆します.

AtCoder Beginner Contest 267 Ex問題 Odd Sum

AtCoder ABC267 ABC Ex問題 600 pts

問題

atcoder.jp

提出解答

atcoder.jp

問題の概要

N 個の整数 A_1, \dots, A_N の中から奇数個選び, 和を M にする選び方の数を求めよ.

制約

  • 1 \leq N \leq 10^5
  • 1 \leq M \leq 10^6
  • 1 \leq A_i \leq 10

解法

制約から, \displaystyle A_{{\rm sum}}:=\sum_{i=1}^N A_i \leq 10^6 である.

まず, A_{{\rm sum}} \lt M ならば, 答えは明らかに 0 なので, 以降は M \leq A_{{\rm sum}} であると仮定する.

このとき, A_{\rm sum} と M の値によって, 次のように K, R を定める.

  • M \leq A_{\rm sum}-M のとき
    • K=M, R=1
  • M \gt A_{\rm sum}-M のとき
    • K=A_{\rm sum}-M
    • N が偶数ならば H=1, N が奇数ならば H=0

このとき, A_1, \dots, A_N の中から和が K になり, 選んだ個数を 2 で割った余りが H であるような選び方の数を求めることにする.

沢山の整数の中からいくつか選んで決まった和になる選び方を求めるので, 多項式を利用する. N 個の2変数多項式 F_i(X,Y) を

F_i(X,Y):=1+X^{A_i} Y

とする. このとき, G:=F_1 \times \dots \times F_N としたとき, G の X^m Y^r の係数は A_1, \dots, A_N の中から r 個選んで和が m になる選び方の数に対応する.

今回求めるべきは選んだ個数を 2 で割った余りに関するので, (Y^0, Y^2, Y^4, \dots), (Y^1, Y^3, Y^5, \dots) をそれぞれ同一視したい. これは Y^2 を 1 とみなすことと同等である.

これを可能にするには, Y^2-1 で割った剰余を考えれば良い. G を Y の多項式と見なし, (Y^2-1) で割った商を Q(X,Y), 余りを R(X,Y) をする.

ここで, R(X,Y) は Y に関して高々1次式なので, R(X,Y)=S(X)+T(X) Y となる X に関する X に関する多項式 S,T が存在する.

これらは

G(X,Y)=Q(X,Y)(Y^2-1)+(S(X)+T(X)Y) \cdots (*)

を満たす. このとき, S が偶数個選ぶことに, T が奇数個選ぶことに対応する多項式である.

よって, この S,T を求められれば良い. (*) に Y=-1, Y=1 をそれぞれ代入すると, Y^2-1=0 になることから,

  • G(X,1)=S(X)+T(X)
  • G(X,-1)=S(X)-T(X)

となることから, これらを解いて.

S(X)=\dfrac{G(X,1)+G(X,-1)}{2}, \quad T(X)=\dfrac{G(X,1)-G(X,-1)}{2}

となる. よって, G(X, \beta) を求めることを目標にする.

G(X,\beta) を高速に求める方法について,

\displaystyle \log(1+\beta X^a)=\sum_{n=1}^{\infty} (-1)^{n-1} \cdot \dfrac{(\beta X^a)^n}{n!}

であるから,

\displaystyle \begin{align} G(X,\beta) &=\exp \left (\log G(X,\beta) \right) \\ &=\exp \left (\log \prod_{i=1}^N F_i(X,\beta) \right) \\ &=\exp \left(\sum_{i=1}^N \log F_i(X,\beta) \right) \\ &=\exp\left(\sum_{i=1}^N \sum_{n=1}^{\infty} (-1)^{n-1} \cdot \dfrac{(\beta X^a )^n}{n!} \right) \end{align}

となる. これの高速化の方法は以下の Library Checker の提出から見てもらうことにして,

judge.yosupo.jp

これにより, S,T のうち, 必要な方の多項式の M 乗の係数をみることにより, 時間計算量 O(N+K \log K) で求めることができる. なお, K の定め方から, K \leq A_{{\rm sum}}/2 \leq 5 \times 10^5 であり, 4 sec の猶予が与えられているので, 間に合う.