ABC388 D問題（Coming of Age Celebration）を解く

AtCoder が提供しているABC（AtCoder Beginner Contest）388 D問題をC++とPythonで解いてみました。ABC388は、2025年1月11日21:00に実施されました。

AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。

D問題 Coming of Age Celebration（Difficulty : 659）

問題の詳細は、リンク先をご覧ください。

ABC388 D問題 Coming of Age Celebration

いもす法によって、先輩から貰う石の個数を求めます。AtCoder Problems による Difficulty は 659 でした。

解答案

宇宙人の石は、以下のような動きをします。

• $i$ 番目の宇宙人は、$A_i$ 個の石を持っている。
• $i$ 年目以降に $i+1, i+2, \cdots, N-1$ 番目の宇宙人に石を一つずつ渡す。
  • このため、$i$ 番目の宇宙人の石は $N-i-1$ 個減る。
  • もらう石は、いもす法によって計算できます。
    • いもす法については、タグを設定しました。そちらの記事も参考にしてください。
    • 典型的な問題は、少し古いですが ABC014 C問題（解説記事）となります。

C++ プログラム例（ABC388D）

プログラムを補足します。

• 先輩からもらう石の数を配列 t で管理します。
  • 石の数 a[i] に、もらう石を加えるために先に t を更新します（17行目）。
  • もらった石 t[i] を持っている石の数 a[i] に加えます（19行目）。
• 配る石の数をいもす法で処理します。
  • 半開区間 [left, right) に値を加えます。このため区間の端である t[left] の値を増やして、t[right] の値を減らします（いもす法）。
• 最後に a[i] から配る石の個数を引いて更新します（25行目）。これが $N$ 年後に持っている石の個数となります。

以下が、C++プログラムです。

#include <bits/stdc++.h>
using namespace std;

typedef long long int ll;
int main()
{
	int n;
	cin >> n;
	vector<int> a(n);
	for (int i = 0; i < n; ++i) {
		cin >> a[i];
	}

	vector<int> t(n + 1);
	for (int i = 0; i < n; ++i) {
		if (i > 0) {
			t[i] += t[i - 1];
		}
		a[i] += t[i];
		int left, right;
		left = i + 1;
		right = min(n, i + a[i] + 1);
		++t[left];
		--t[right];
		a[i] = max(a[i] - (n - 1 - i), 0);
	}

	for (int i = 0; i < n; ++i) {
		cout << a[i] << " \n"[i == n - 1];
	}

	return 0;
}

AC（Accepted＝正しいプログラム）と判定されました。

Python プログラム例（ABC388D）

Python版も基本的な考え方は同じです。以下がプログラムです。

"""AtCoder Beginner Contest 388 D"""
n = int(input())
a = list(map(int, input().split()))

t = [0] * (n + 1)
for i in range(n):
    if i > 0:
        t[i] += t[i - 1]
    a[i] += t[i]
    left = i + 1
    right = min(n, i + a[i] + 1)
    t[left] += 1
    t[right] -= 1
    a[i] = max(a[i] - (n - 1 - i), 0)

print(*a)

こちらも「AC」と判定されました。

最後に

今回のコンテストは、E問題に時間を使いました。結果的にこのD問題は、いもす法を使うことは思いついていましたが時間切れでした。残念でした。

引き続き ABC の問題を紹介していきます。