AtCoder

ABC384 D問題(Repeated Sequence)を解く

AtCoder_ABC384_D

AtCoder が提供しているABC(AtCoder Beginner Contest)384 D問題をC++とPythonで解いてみました。ABC384は、2024年12月14日21:00に実施されました。

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

D問題 Repeated Sequence(Difficulty : 666)

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

ABC384 D問題 Repeated Sequence

累積和を求めて二分探索します。AtCoder Problems による Difficulty は 666 でした。

解答案

周期 $N$ をもつ無限数列を対象とするため、$S$ を $\sum a_i$ で割った余りについて考えればよいことが分かります。

数列 $A$ の後半から開始して、数列の前半で終わる場合があるため、数列 $A$ の後ろに $A$ の要素を加えて考察します。

C++ プログラム例(ABC384D)

数列 $A$ を2回繰り返した数列の累積和を求めます(17ー20行目)。累積和の配列は、昇順に並んでいるため、目的の値が存在するか二分探索で求めることができます(25ー28行目)。

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

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

typedef long long int ll;

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

	vector<ll> sum_a(2 * n + 1, 0);
	for (int i = 0; i < 2 * n; ++i) {
		sum_a[i + 1] = sum_a[i] + a[i];
	}

	bool result = false;
	s %= sum_a[n];
	for (int i = 0; i < 2 * n; ++i) {
		auto it = lower_bound(sum_a.begin(), sum_a.end(), sum_a[i] + s);
		if ((it != sum_a.end()) && (*it == sum_a[i] + s)) {
			result = true;
		}
	}

	if (result) {
		cout << "Yes" << endl;
	} else {
		cout << "No" << endl;
	}

	return 0;
}

AC(Accepted=正しいプログラム)と判定されました。

Python プログラム例(ABC384D)

Python版も基本的な考え方は同じです。二分探索には、bisect.bisect_left を使いました。以下がプログラムです。

"""AtCoder Beginner Contest 384 D"""
import bisect

n, s = map(int, input().split())
a = list(map(int, input().split()))
for i in range(n):
    a.append(a[i])

sum_a = [0] * (2 * n + 1)
for i in range(2 * n):
    sum_a[i + 1] = sum_a[i] + a[i]

result = False
s %= sum_a[n]
for i in range(2 * n):
    it = bisect.bisect_left(sum_a, sum_a[i] + s)
    if it < 2 * n and sum_a[it] == sum_a[i] + s:
        result = True

print("Yes" if result else "No")

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

最後に

コンテストでは、配列 a を2組格納したため、sum_a[2 * n] の剰余を求めていました(sum_a[n] の剰余を求めるのが正しい)。58テストケースで3件がWAとなっており、コンテスト中にこの不具合に気が付きませんでした。残念でした。

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

COMMENT

メールアドレスが公開されることはありません。 が付いている欄は必須項目です

CAPTCHA