ABC340 E問題（Mancala 2）を解く

AtCoder が提供しているABC（AtCoder Beginner Contest）340 のE問題をC++で解いてみました。ABC340は、2024年2月10日21:00に実施されました。

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

E問題 Mancala 2（Difficulty : 1157）

問題はリンク先をご覧ください。

ABC340 E問題 Mancala 2

遅延評価セグメント木を使います。AtCoder Problems による Difficulty は 1157 でした。

解答案

この問題は、遅延評価ができるセグメント木を使うだけで解くことができます。具体的には、以下となります。

• c = a[b[i]] とする。
• c / n が1以上なら、c を 0 にした後に、すべての a の要素に a[b[i]] / n を加える。
• 次に c = c % n とする。
  • c が 0 なら何もしない。
  • a の添え字 b[i] + 1 から c 個の要素に 1 を加える。
  • 添え字が n 以上になれば、0 に戻して 1 を加える。

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

AtCoder Library には、遅延評価セグメント木 lazy_segtree が用意されているため、これを使います。上で整理した内容に従い、lazy_segtree のメソッドを呼び出すだけです。

以下が、C＋＋プログラムとなります。

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

typedef long long int ll;

ll op(ll a,ll b)
{
	return a + b;
}

ll e()
{
	return 0LL;
}

int main()
{
	ll n, m;
	cin >> n >> m;

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

	lazy_segtree<ll, op, e, ll, op, op, e> st(a);
	for (int i = 0; i < m; ++i) {
		ll c = st.get(b[i]);
		st.set(b[i], 0);
		if (c / n > 0) {
			st.apply(0, n, c / n);
		}
		c = c % n;
		if (c == 0) {
			continue;
		} else if (b[i] + c < n) {
			st.apply(b[i] + 1, b[i] + c + 1, 1);
		} else {
			st.apply(b[i] + 1, n, 1);
			st.apply(0, b[i] + c - n + 1, 1);
		}
	}

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

	return 0;
}

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

Python の遅延評価セグメント木について

Python は、遅延評価セグメント木を標準ライブラリとしては用意していません。このブログの Python プログラムは、なるべく言語標準のライブラリで実装したいと考えています。このような事情で、この記事では、Python に移植したプログラムの紹介を省略します。

最後に

遅延評価セグメント木を使う問題としては初歩的な内容でしょうか。コンテストでは解くことができませんでしたが、解説を読み、遅延評価セグメント木について理解できました。

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