AtCoder Beginner Contest 253 E問題 Distance Sequence

問題

atcoder.jp

提出解答

atcoder.jp

問題の概要

次を満たす長さ $N$ の整数列 $A=(A_1, \dots, A_N)$ の個数を求めよ.

$1 \leq A_i \leq M$
$\left|A_i-A_{i+1} \right| \geq K$

制約

$2 \leq N \leq 1000$
$1 \leq M \leq 5000$
$0 \leq K \leq M-1$

解法

動的計画法で解く. ${\rm DP}[i,j]~(1 \leq i \leq N, 1 \leq j \leq M)$ を次のように定義する.

長さが $i$ である条件を満たすような列のうち, 末項が $j$ であるような数列の数.

このとき, 最終解答は

$\displaystyle \sum_{j=1}^M {\rm DP}[i,j]$

である.

ベースケースは $i=1$ のときであり, 任意の $1 \leq j \leq M$ に対して, ${\rm DP}[1,j]=1$ である.

遷移式を考える. 第 $i$ 項として $j$ が可能かどうかは第 $(i-1)$ 項目だけにより決定づけられることに注意すると,

$\displaystyle {\rm DP}[i,j]=\sum_{\substack{1 \leq k \leq M \\ \left| k-j \right| \geq K}} {\rm DP}[i-1,k]$

である. これをそのまま実装してしまうと, 各 ${\rm DP}[i,j]$ を求めるために $O(M)$ 時間かかり, 動的計画法の要素数は $NM$ なので, 全体で $O(NM^2)$ 時間となり間に合わない.

そこで, 更新式を計算して高速に処理できるような形になる.

$\displaystyle S_{i,j}:=\begin{cases} \displaystyle \sum_{k=1}^j {\rm DP}[i,k] & (j \geq 1) \\ 0 & (j=0) \end{cases}$

とおくと,

$\displaystyle \begin{align} {\rm DP}[i,j] &=\sum_{\substack{1 \leq k \leq M \\ \left| k-j \right| \geq K}} {\rm DP}[i-1,k] \\ &=\sum_{k=1}^M {\rm DP}[i-1,k]-\sum_{\substack{1 \leq k \leq M \\ \left| k-j \right| \lt K}} {\rm DP}[i-1,k] \\ &=\sum_{k=1}^M {\rm DP}[i-1,k]-\begin{cases} 0 & (K=0) \\ \displaystyle \sum_{\substack{1 \leq k \leq M \\ j-(K-1) \leq k \leq j+(K-1)}} {\rm DP}[i-1,k] & (K \geq 1) \end{cases} \\ &=\sum_{k=1}^M {\rm DP}[i-1,k]-\begin{cases} 0 & (K=0) \\ \displaystyle \sum_{k=1}^{\min (M, j+(K-1))} {\rm DP}[i-1,k]-\sum_{k=1}^{\max (1, j-(K-1))-1} {\rm DP}[i-1,k] & (K \geq 1) \end{cases} \\ &=S_{i,M}-\begin{cases} 0 & (K=0) \\ S_{i-1, \min (M, j+(K-1))}-S_{i-1, \max(1, j-(K-1))-1} & (K \geq 1) \end{cases} \end{align}$

となる.

$i=1,2, \dots, N$ の順に ${\rm DP}[i,j]$ を求めることにすると, 各 $i (\geq 2)$ に対して, $S_{i-1, 1}, \dots, S_{i-1, M}$ を前計算で求めておくと, 前計算 $O(M)$ 時間, 各 $j$ に対して, ${\rm DP}[i, j]$ を求めるために $O(1)$ 時間で処理できる.