実は2～3年前、競プロを1年弱くらいやっていたのだけど、緑に入ったくらいで一度やめてしまっていた。ところが、最近始めたバイトは競プロっぽい問題の作問*1をする仕事なので、久しぶりにAtCoderをやってみようと思って再開した。以前はC++で書いていたけど、せっかくだからPythonで書いてみようと思って、最近はもっぱらリハビリとしてABCのD問題以下を適当に解いてます。
今日解いていた問題はこちら： atcoder.jp
AtCoder Problems のリコメンドで出てきたので解くことにした。これ、かなり便利なのでいつも使ってます。
問題の解説は省略。まあいい感じにDPと累積和を組み合わせる。DPを書いたのが久しぶりだったので実装にかなり時間がかかってしまった。
そこでできたコードがこちら：
n = int(input()) a = list(map(int, input().split())) b = list(map(int, input().split())) dp = [[0] * 3001 for i in range(n)] dp[0][a[0]] = 1 for j in range(3000): dp[0][j + 1] += dp[0][j] for i in range(n - 1): for j in range(a[i], b[i] + 1): dp[i + 1][j] = dp[i][j] for j in range(3000): dp[i + 1][j + 1] += dp[i + 1][j] print(sum(dp[n - 1][j] for j in range(a[n - 1], b[n - 1] + 1)) % 998244353)
結果はTLE。う～んと思って、C++で書き直してみる。ここで、C++のintは多長倍整数型でないので累積和の処理の中でmodを取らないといけないことを思い出す。
#include <bits/stdc++.h> using namespace std; int main() { int n; cin >> n; int a[n], b[n]; for (int i = 0; i < n; i++) cin >> a[i]; for (int i = 0; i < n; i++) cin >> b[i]; int dp[n][3001]; for (int i = 0; i < n; i++) for (int j = 0; j < 3001; j++) if (i == 0 && a[0] <= j && j <= b[0]) { dp[i][j] = 1; } else { dp[i][j] = 0; } for (int i = 0; i < n - 1; i++) { for (int j = a[i]; j <= b[i]; j++) dp[i + 1][j] = dp[i][j]; for (int j = 0; j < 3000; j++) dp[i + 1][j + 1] = (dp[i + 1][j + 1] + dp[i + 1][j]) % 998244353; } int sum = 0; for (int j = a[n - 1]; j <= b[n - 1]; j++) sum = (sum + dp[n - 1][j]) % 998244353; cout << sum << endl; }
今度は通った。あ～多長倍整数型は桁が大きくなったら遅くなるんだな～と思って途中でmodを取るやつをPythonで書いて提出。でもTLE。えーお手上げだ～と思って「競プロ Python TLE」で検索。
この記事によると何やらPyPy3で提出したら通るらしい。やってみた。通った。なんじゃこりゃ。
Pythonは書くコード量が少なくなるのはいいけど、こういう定数倍高速化の手間が必要ならC++のほうが考えることが少なくてラクかもしれないなと思ったり。これからもちょくちょく競プロの記事を書こうと思う。次は久々にRatedコンテストに出ようかと思ってます
