AtCoder が提供しているABC(AtCoder Beginner Contest)336 のD問題をC++とPythonで解いてみました。ABC336は、2024年1月14日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 Pyramid(Difficulty : 991)
問題はリンク先をご覧ください。
前からと後ろから計算して、統合して解を求めます。AtCoder Problems による Difficulty は 991 でした。
解答案
入力例から、左から考えた山の高さ、右から考えた山の高さ、2つの値の小さい方の値を計算します。
入力例1
数列A | 2 | 2 | 3 | 1 | 1 |
左から考えた山の高さ | 1 | 2 | 3 | 1 | 1 |
右から考えた山の高さ | 2 | 2 | 2 | 1 | 1 |
min(左から考えた山の高さ, 右から考えた山の高さ) | 1 | 2 | 2 | 1 | 1 |
数列A | 1 | 2 | 3 | 4 | 5 |
左から考えた山の高さ | 1 | 2 | 3 | 4 | 5 |
右から考えた山の高さ | 1 | 2 | 3 | 2 | 1 |
min(左から考えた山の高さ, 右から考えた山の高さ) | 1 | 2 | 3 | 2 | 1 |
入力例1と2は、どちらも min(左から考えた山の高さ, 右から考えた山の高さ) の最大値が解答になっています。
左から考えた山の高さ dp_u は、以下の漸化式で求めることができます。
- dp_u[0] = 1
- dp_u[i+1] = min(dp_u[i]+1, a[i+1]) i=0 から n-1 の順で更新
右から考えた山の高さ dp_d は、以下の漸化式で求めることができます。
- dp_d[n-1] = 1
- dp_d[i-1] = min(dp_d[i]+1, a[i-1]) i=n-1 から 1 の順で更新
C++ プログラム例(ABC336D)
上記の考察の漸化式をプログラムにしました。左から考えた山の高さと右から考えた山の高さの小さい方の値の最大値を求めます(26行目)。
以下が、C++プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
int main()
{
int n;
cin >> n;
vector<int> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
vector<int> dp_u(n);
dp_u[0] = 1;
for (int i = 0; i < n - 1; ++i) {
dp_u[i + 1] = min(dp_u[i] + 1, a[i + 1]);
}
vector<int> dp_d(n);
dp_d[n - 1] = 1;
for (int i = n - 1; i >= 1; --i) {
dp_d[i - 1] = min(dp_d[i] + 1, a[i - 1]);
}
int result = 1;
for (int i = 0; i < n; ++i) {
result = max(result, min(dp_u[i], dp_d[i]));
}
cout << result << endl;
return 0;
}
AC(Accepted=正しいプログラム)と判定されました。
Python プログラム例(ABC336D)
Python 版も基本的な考え方は同じです。以下となります。
"""AtCoder Beginner Contest 336 D"""
n = int(input())
a = list(map(int, input().split()))
dp_u = [0] * n
dp_u[0] = 1
for i in range(n - 1):
dp_u[i + 1] = min(dp_u[i] + 1, a[i + 1])
dp_d = [0] * n
dp_d[n - 1] = 1
for i in range(n - 1, 0, -1):
dp_d[i - 1] = min(dp_d[i] + 1, a[i - 1])
result = 1
for i in range(n):
result = max(result, min(dp_u[i], dp_d[i]))
print(result)
こちらも「AC」と判定されました。
最後に
コンテストで、AからC問題のACを 9分17秒で解くことができました。一方、D問題は一回でACが取れましたが、時間は89分23秒までかかりました。
提出したプログラムは、AC判定でしたが、ごちゃごちゃしたコードになっています。すっきりと考えることができなかったため、80分もかかったと考えています。
コンテストで提出したコードを反省のために貼っておきます。
#include <bits/stdc++.h>
using namespace std;
int main()
{
int n;
cin >> n;
vector<int> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
vector<int> u(n, 0);
int max_u = 1;
u[0] = 1;
for (int i = 0; i < n - 1; ++i) {
if ((a[i] >= max_u)&&(a[i + 1] >= max_u + 1)) {
u[i + 1] = max_u + 1;
++max_u;
} else if ((a[i] >= max_u)&&(a[i + 1] < max_u + 1)) {
u[i + 1] = min(max_u, a[i + 1]);
max_u = u[i + 1];
} else if ((a[i] < max_u)&&(a[i + 1] >= max_u + 1)) {
u[i + 1] = a[i] + 1;
max_u = u[i + 1];
} else {
u[i + 1] = min(a[i] + 1, a[i + 1]);
max_u = u[i + 1];
}
}
vector<int> d(n, 0);
int max_d = 1;
d[n - 1] = 1;
for (int i = n - 1; i >= 1; --i) {
if ((a[i] >= max_d)&&(a[i - 1] >= max_d + 1)) {
d[i - 1] = max_d + 1;
++max_d;
} else if ((a[i] >= max_d)&&(a[i - 1] < max_d + 1)) {
d[i - 1] = min(max_d, a[i - 1]);
max_d = d[i - 1];
} else if ((a[i] < max_d)&&(a[i - 1] >= max_d + 1)) {
d[i - 1] = a[i] + 1;
max_d = d[i - 1];
} else {
d[i + 1] = min(a[i] + 1, a[i - 1]);
max_d = d[i - 1];
}
}
int result = 1;
for (int i = 0; i < n; ++i) {
result = max(result, min(u[i], d[i]));
}
cout << result << endl;
return 0;
}