Kazun の競プロ記録

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

AtCoder Beginner Contest 262 D問題 I Hate Non-integer Number

AtCoder ABC ABC262 C問題 300 pts

問題

提出解答

問題の概要

長さ $N$ の非負整数列 $A=(a_1, \dots, a_N)$ が与えられる. $A$ から $1$ 個以上の項を選ぶ方法のうち, それらの平均が整数であるような選び方の数を求めよ.

制約

$1 \leq N \leq 100$
$1 \leq a_i \leq 10^9$

解法

以下は全て同値である. ただし, $k \geq 1$ とする.

$k$ 個の整数 $x_1, \dots, x_k$ の平均が整数
$k$ 個の整数 $x_1, \dots, x_k$ の総和が $k$ の倍数

従って, 次の小問題を $k=1,2, \dots, N$ に対して解くことになる.

$A$ の項からちょうど $k$ 個選ぶ方法のうち, 和が $k$ の倍数であるような選び方の数を求めよ.

動的計画法で解くことにする. ${\rm DP}_k[i,p,m]$ で第 $i$ 項目まで見たとき, それまでに $p$ 個選んで, 和が $k$ で割ると $m$ になるような選び方の数とする.

ベースケースは $i=0$ のときで

${\rm DP}_k[0,p,m]=\begin{cases} 1 & (p=0, m=0) \\ 0 & ({\rm otherwise}) \end{cases}$

である.

更新式は第 $i$ 項目を採用するかしないかで場合分けすることにより,

${\rm DP}_k[i,p,m]=\begin{cases} {\rm DP}_k[i-1,0,m] & (p=0) \\ {\rm DP}_k[i-1,p-1,(m-a_i) \pmod{k}]+{\rm DP}_k[i-1,p,m] & (p \geq 1) \end{cases}$

となる.

このとき, 小問題の答えは ${\rm DP}_k[N,k,0]$ である. 従って, 全体の答えは

$\displaystyle \sum_{k=1}^N {\rm DP}_k[N,k,0]$

である. 更新式が $O(1)$ 時間で, 要素数が $O(N^4)$ であるから, この問題を $O(N^4)$ で解くことができた.