ABC262 Dの解法メモ【Ruby】

入水目指して過去挑戦して解けてないAtCoder Problems上でDifficultyが水色の問題を解いていく。

ABC262 D: I Hate Non-integer Number概要

長さNの整数列の部分列のうち、要素の平均が整数になる部分列の数を求める問題。

解法イメージ

$ 2^{100} = 2^{10 \times 10} = 1024^{10} \fallingdotseq 10^{3 \times 10} $なので全パターンを数え上げるのは制限時間内には間に合わない。
$ i $個の整数の平均が整数になるということは和を$ i $で割った時の余りが0になるので、各$ i $に対して各整数の剰余を管理すればよさそう。
$ 1 \le N \le 100 $であり部分列の和を計算するのでDPが使えそうなので検討してみる。

各$ i $に対して部分和のDPを考えるので、$ j $番目までの整数から$ k $個の整数を選択した場合として考える。
今回は$ i $で割った時の余りが0になる個数を計算したいので、$ k $個の和の剰余($ mod\ i $)が$ l $となる部分列の個数を管理してあげると、$ l^{'} \equiv l - A_j \ (mod\ i)$とすると

dp[j][k][l] = \left\{
\begin{array}{ll}
dp[j-1][k][l] & (k = 0) \\
dp[j-1][k][l] + dp[j-1][k-1][l^{'}] & (k \gt 0) \\
dp[j-1][k-1][l^{'}] & (k = j)
\end{array}
\right. 

で計算できる。

求めたいのは$ dp[N][i][0] $なので、これらを各$ i $について求めた合計値を出力すればよい。

実装例

n = gets.to_i
arr = gets.split.map(&:to_i)
mod = 998244353

c = 0
# dp[0]は固定で以降は毎回再計算されるのでiのループで初期化をする必要はない
dp = Array.new(n + 1) { Array.new(n + 1) { Array.new(n) { 0 }}}
dp[0][0][0] = 1
(1..n).each do |i| # i個の平均が整数になるケース
  (1..n).each do |j| # j番目の要素までで
    a = arr[j - 1]
    ([i, j].min+1).times do |k|| # k個の要素を足したときに
      i.times do |l|   # iで割った時の余りがlになるもの
        # dp[j-1][-1]は常に0の配列になるので問題ない
        dp[j][k][l] = (dp[j-1][k][l] + dp[j-1][k-1][(l-a)%i]) % mod
      end
    end
  end
  c = (c + dp[n][i][0]) % mod
end
puts c