Typical DP Contest F - 準急
半日かかった。dp苦手すぎる。
F - 準急
問題
N個の駅のある路線を準急が走る。以下の条件を満たす停車駅の組み合わせの数を求めよ。
- 1個目、N個目の駅に止まる。
- 連続してK個の駅に止まることはない。
解法
まず、最後にi
番目の駅に停車しそれが連続停車j
駅目であるような停車駅の組み合わせdp[i][j]
を計算することを考える。
i-2
駅目かそれ以前の駅に最後に止まる組み合わせ全ての数がdp[i][1]
であり、j >= 2
に関してはdp[i][j] = dp[i-1][j-1]
である。
単純にすると$O(N^2K)$と間に合わない。
dp[i][1]
を計算する部分で累積和を用い、dp[i][1] = acc[i-2]
とすると$O(NK)$。
dp[i][j] = dp[i-1][j-1] = ... = dp[i-j+1][1] = acc[i-j-1]
であることを用いて、dp tableから連続で何駅停車したかの情報を落とせば$O(N)$。ただし負の引数のacc
が何を表すために用いられているかに注意。
実装
$O(n)$解
#include <iostream>
#include <vector>
#include <cassert>
#define repeat_from(i,m,n) for (int i = (m); (i) < (n); ++(i))
#define repeat(i,n) repeat_from(i,0,n)
typedef long long ll;
using namespace std;
constexpr ll mod = 1000000007;
int main() {
int n, k; cin >> n >> k;
vector<ll> dp(n);
vector<ll> acc(n);
dp[0] = 1;
acc[0] = 1;
repeat_from (x,1,n) {
ll ex = x-k+1 == 0 ? 1 : 0 <= x-k-1 ? acc[x-k-1] : 0;
dp[x] = (acc[x-1] - ex + mod) % mod;
acc[x] = (acc[x-1] + dp[x]) % mod;
}
cout << dp[n-1] << endl;
return 0;
}
$O(n^2k)$解
直接実装するだけの能力がないので、確実に正しい答えを返すものを少しずつ修正していった。その過程でできたもの。
#include <iostream>
#include <vector>
#include <cassert>
#define repeat_from(i,m,n) for (int i = (m); (i) < (n); ++(i))
#define repeat(i,n) repeat_from(i,0,n)
typedef long long ll;
using namespace std;
constexpr ll mod = 1000000007;
int main() {
int n, k; cin >> n >> k;
vector<vector<ll> > dp(n, vector<ll>(k-1));
dp[0][0] = 1;
repeat_from (x,1,n) {
repeat (y,x-1) {
repeat (z,k-1) {
dp[x][0] += dp[y][z];
dp[x][0] %= mod;
}
}
repeat_from (y,1,k-1) {
dp[x][y] += dp[x-1][y-1];
dp[x][y] %= mod;
}
}
ll result = 0;
repeat (x,k-1) {
result += dp[n-1][x];
result %= mod;
}
cout << result << endl;
return 0;
}