AtCoder が提供しているABC(AtCoder Beginner Contest)359 E問題をC++とPythonで解いてみました。ABC359は、2024年6月22日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Water Tank(Difficulty : 1275)
問題の詳細は、リンク先をご覧ください。
スタックを使って、それより前にある高い板を求めます。AtCoder Problems による Difficulty は 1275 でした。
解答案
それより前(左)にある板で、自身の $H_i$ より高い板とその位置を見つけることができれば解くことができます。
このためにスタックを使います。スタックに (板の高さ, 位置) を格納します。
C++ プログラム例(ABC359E)
次の工夫をします。
- 左端の壁として高さ無限大の板を設定します(12、19行目)。
- 解答となる配列 result には、あふれる直前の水の量を格納します。あふれた時間を出力する必要があるため、出力する際に1を加えます(31行目)。
- 直前にある自身より高い板を求めます。その過程で pop しながら探すため、時間的に間に合います(21ー28行目)。
以下が、C++プログラムです。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
#define INF (1LL << 60)
int main()
{
int n;
cin >> n;
vector<ll> h(n + 1);
h[0] = INF;
for (int i = 1; i <= n; ++i) {
cin >> h[i];
}
vector<pair<ll, int>> st;
vector<ll> result(n + 1);
st.push_back(make_pair(INF, 0));
result[0] = 0;
for (int i = 1; i <= n; ++i) {
while (st.back().first <= h[i]) {
st.pop_back();
}
int pre_pos = st.back().second;
result[i] = (i - pre_pos) * h[i] + result[pre_pos];
st.push_back(make_pair(h[i], i));
}
for (int i = 1; i <= n; ++i) {
cout << result[i] + 1 << " \n"[i == n];
}
return 0;
}
AC(Accepted=正しいプログラム)と判定されました。
Python プログラム例(ABC359E)
Python版も基本的な考え方は同じです。以下がプログラムです。
"""AtCoder Beginner Contest 359 E"""
INF = 1 << 60
n = int(input())
h = list(map(int, input().split()))
h.insert(0, 0)
st = []
result = [0] * (n + 1)
st.append((INF, 0))
for i in range(1, n + 1):
while st[-1][0] <= h[i]:
st.pop()
pre_pos = st[-1][1]
result[i] = (i - pre_pos) * h[i] + result[pre_pos]
st.append((h[i], i))
for i in range(1, n + 1):
print(result[i] + 1, end=" ")
print()
こちらも「AC」と判定されました。
最後に
近い考え方で解く問題として、ALDS1 3_D問題(解説記事)をすでに紹介していました。
コンテスト中は、setコンテナで解こうとして悪戦苦闘して解けませんでした。スタックを使って実装したら、あっさりと解けました。勉強になりました。
引き続き ABC の問題を紹介していきます。