AtCoder が提供しているABC(AtCoder Beginner Contest)347 のE問題をC++とPythonで解いてみました。ABC347は、2024年3月30日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Set Add Query(Difficulty : 1166)
問題はリンク先をご覧ください。
それぞれの整数について、集合に含まれている期間を記録して、各要素の値を累積和で求めます。AtCoder Problems による Difficulty は 1166 でした。
解答案
$1 \leqq x_i \leqq N$ について、最初(奇数回目)に指定されたときに集合に加えて、2回目(偶数回目)に指定されたときに集合から削除します。集合に加えると同時にそれぞれの数字が出現したクエリ順を記憶しておきます。
集合の要素を加えていく計算は累積和を求めて、[奇数回目, 偶数回目) の範囲の和を集計して、それぞれの整数の値とします。
C++ プログラム例(ABC347E)
プログラムの補足です。
- 要素のインデックスは0からカウントします(13行目)。
- set コンテナ s に数字を加えるか、削除します(20ー24行目)。
- コンテナ s の要素数から累積和を計算します(25行目)。
- もし、奇数回しか指定されなかった場合は、末尾+1の q をクエリ順に加えます。
- [奇数回目, 偶数回目) の和を出力します(32ー36行目)。
以下が、C++プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
int main()
{
int n, q;
cin >> n >> q;
vector<int> x(q);
for (int i = 0; i < q; ++i) {
cin >> x[i];
--x[i];
}
set<int> s;
vector<vector<int>> v(n);
vector<ll> sum(q + 1, 0);
for (int i = 0; i < q; ++i) {
v[x[i]].push_back(i);
if (s.find(x[i]) == s.end()) {
s.insert(x[i]);
} else {
s.erase(x[i]);
}
sum[i + 1] = sum[i] + s.size();
}
for (int i = 0; i < n; ++i) {
if (v[i].size() % 2 != 0) {
v[i].push_back(q);
}
ll t_sum = 0;
for (int j = 0; j < v[i].size() / 2; ++j) {
t_sum += sum[v[i][j * 2 + 1]] - sum[v[i][j * 2]];
}
cout << t_sum << " ";
}
cout << endl;
return 0;
}
AC(Accepted=正しいプログラム)と判定されました。
Python プログラム例(ABC347E)
Python 版も基本的な考え方は同じです。以下となります。
"""AtCoder Beginner Contest 347 E"""
n, q = map(int, input().split())
x = list(map(int, input().split()))
for i in range(q):
x[i] -= 1
s = set()
v = [[] for i in range(n)]
sum = [0] * (q + 1)
for i in range(q):
v[x[i]].append(i)
if x[i] not in s:
s.add(x[i])
else:
s.remove(x[i])
sum[i + 1] = sum[i] + len(s)
result = []
for i in range(n):
if len(v[i]) % 2 != 0:
v[i].append(q)
t_sum = 0
for j in range(len(v[i]) // 2):
t_sum += sum[v[i][j * 2 + 1]] - sum[v[i][j * 2]]
result.append(t_sum)
print(*result)
こちらも「AC」と判定されました。
最後に
コンテスト当日は、D問題で時間を使ってしまい、この問題を解く時間が無くなりました。コンテスト後に解きましたが、set コンテナと累積和を組み合わせる問題で、学びが多い問題でした。
引き続き ABC の問題を紹介していきます。