Typical DP Contest F: 準急 (動的計画法)

問題

F: 準急 - Typical DP Contest | AtCoder

解法

(駅i+1で連続j+1回目の停車をする組み合わせ)
 = (駅iで連続j回目の停車をする組み合わせ)

なので、愚直にやってもO(NK)で解ける。工夫して、

(駅i+1で停車する組み合わせ)
 = (駅iで停車する組み合わせ) + (駅iで通過する組み合わせ)
 - (駅i+1-Kで通過する組み合わせ)

とするとO(N)で解けるようになる。

コード

import java.io.IOException;

public class Main {

	private static final int MOD = 1000000007;

	public static void main(String[] arg) {
		int N = nextInt();
		int K = nextInt();

		// dp[i][0]:=駅iを通過する組み合わせの数
		// dp[i][1]:=駅iで止まる組み合わせの数
		long[][] dp = new long[N + 1][2];
		dp[0][0] = 1;
		dp[1][1] = 1;

		for (int n = 1; n < N; n++) {
			dp[n + 1][1] += dp[n][0] + dp[n][1];
			if (n + 1 - K >= 0) {
				dp[n + 1][1] -= dp[n + 1 - K][0];
			}
			dp[n + 1][1] += MOD * 2;
			dp[n + 1][1] %= MOD;

			dp[n + 1][0] += dp[n][0] + dp[n][1];
			dp[n + 1][0] %= MOD;
		}

		System.out.println(dp[N][1]);

	}

	static int nextInt() {
		int c;
		try {
			c = System.in.read();
			while (c != '-' && (c < '0' || c > '9'))
				c = System.in.read();
			if (c == '-')
				return -nextInt();
			int res = 0;
			while (c >= '0' && c <= '9') {
				res = res * 10 + c - '0';
				c = System.in.read();
			}
			return res;
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		return -1;
	}

}