AtCoder が提供しているABC(AtCoder Beginner Contest)285 のE問題をC++とPythonで解いてみました。ABC285は、2023年1月15日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Work or Rest(Difficulty : 1466)
問題はリンク先をご覧ください。
DPで解けます。漸化式はひと工夫が必要です。AtCoder Problems による Difficulty は 1466 でした。
解答案
1日ある休みは初日に固定します。これは、このような問題のコツです。$i$ 日連続で働いた場合のトータルの生産量 $P_i$ を求めておきます。$A_i$ のインデックスは0からカウントしています。最後の漸化式の割り算は切り捨てです。
- $P_1 = A_0$
- $P_2 = A_0 + A_0$
- $P_3 = A_0 + A_0 + A_1$
- $P_N = P_{N-1} + A_{(N-1) / 2}$
dp[i][j] を i 日後、この時点で j 日連続で働いている場合の生産量の最大値とします。ただし、j が 0 でない場合は、生産量を確定させません。漸化式は以下となります。初日を0日目とします。
- dp[i][0] = dp[i – j][0] + p[j – 1]
- i 日目を休みだとします。j 日前が休みだとして、生産量を加えます。
- dp[i][j] = dp[i – 1][j – 1]
- 働き続ける場合は、前日の値を引き継ぎます。
$N-1$ 日目まで更新出来た後は、$N$日目は休みとなるため、以下の式で生産量を確定します。
- dp[n – 1][j] + p[j]
C++ プログラム例(ABC285E)
上で考察した漸化式をプログラムに変換しました。以下は、C++ プログラムです。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
#define INF (1ULL << 60)
int main()
{
int n;
cin >> n;
vector<ll> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
vector<ll> p(n, 0);
for (int i = 0; i < n - 1; ++i) {
p[i + 1] += p[i] + a[i / 2];
}
vector<vector<ll>> dp(n, vector<ll>(n, -INF));
dp[0][0] = 0;
for (int i = 1; i < n; ++i) {
for (int j = 1; j < n; ++j) {
if ((i - j >= 0)&&(dp[i - j][0] >= 0)) {
dp[i][0] = max(dp[i][0], dp[i - j][0] + p[j - 1]);
}
dp[i][j] = dp[i - 1][j - 1];
}
}
ll result = dp[n - 1][0];
for (int i = 1; i < n; ++i) {
result = max(result, dp[n - 1][i] + p[i]);
}
cout << result << endl;
return 0;
}
AC(Accepted=正しいプログラム)と判定されました。
Python プログラム例(ABC285E)
Python 版も、C++ のプログラムをそのまま移植しました。
以下のプログラムは、「Python (CPython 3.11.4)」では、TLE判定となり、「Python (PyPy 3.10-v7.3.12)」で、AC判定となりました。
"""AtCoder Beginner Contest 285 E"""
INF = 1 << 60
n = int(input())
a = list(map(int, input().split()))
p = [0] * n
for i in range(n - 1):
p[i + 1] += p[i] + a[i // 2]
dp = [[-INF for _ in range(n)] for _ in range(n)]
dp[0][0] = 0
for i in range(1, n):
for j in range(1, n):
if i - j >= 0 and dp[i - j][0] >= 0:
dp[i][0] = max(dp[i][0], dp[i - j][0] + p[j - 1])
dp[i][j] = dp[i - 1][j - 1]
result = dp[n - 1][0]
for i in range(1, n):
result = max(result, dp[n - 1][i] + p[i])
print(result)
最後に
この問題のDPは、漸化式で値が確定しないこともあり難しく感じました。公式解説などを参考にプログラムを書いて記事にしました。記事にすることで理解度が上がることを期待しています。まぁすぐに忘れてしまいますが。
引き続き ABC の問題を紹介していきます。