AtCoder が提供しているABC(AtCoder Beginner Contest)384 D問題をC++とPythonで解いてみました。ABC384は、2024年12月14日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 Repeated Sequence(Difficulty : 666)
問題の詳細は、リンク先をご覧ください。
累積和を求めて二分探索します。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 の問題を紹介していきます。