AtCoder が提供しているABC（AtCoder Beginner Contest）278 のD問題をC++とPythonで解いてみました。ABC278は、2022年11月19日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 All Assign Point Add（Difficulty : 559）
問題はリンク先をご覧ください。
ABC278 D問題 All Assign Point Add
単純に実装すると、時間が間に合いません。一工夫が必要な問題です。AtCoder Problems による Difficulty は、559 でした。
解答案
問題を整理します。
- n と 数列 A を読み込む。
- q を読み、以下を q 回繰り返す。
- t を読み込む。
- t = 1 の場合、x を読み、配列 A のすべての要素を x にする。
- t = 2 の場合、i と x を読み、A[i] に x を加える。
- t = 3 の場合、i を読み、A[i] を出力する。
- t を読み込む。
読み込んだクエリが t = 1 の場合、数列を読み込んだ値で初期化します。このクエリをどのように処理するかが、この問題を解くポイントとなります。
C++ プログラム例（ABC278D）
最初のプログラム
問題に記載してあることを素直に実装します。数列 Ai も 問題 Qi も 1 から始まるようにしました。プログラムとしては、それほど難しくありません。B問題の難易度でしょうか。
#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long int ull;
int main()
{
int n;
cin >> n;
vector<ull> a(n + 1);
for (int i = 1; i <= n; ++i) {
cin >> a[i];
}
int q;
cin >> q;
for (int i = 1; i <= q; ++i) {
int t, iq, xq;
cin >> t;
if (t == 1) {
cin >> xq;
for (int j = 1; j <= n; ++j) {
a[j] = xq;
}
} else if (t == 2) {
cin >> iq >> xq;
a[iq] += xq;
} else {
cin >> iq;
cout << a[iq] << endl;
}
}
return 0;
}問題に示されている3組の入力例に対して正しい出力をしました。ただし、AtCoder サーバに提出すると、TLE（Time Limit Exceeded）と判定されます。
プログラムを見直すと、t = 1 のクエリで、数列全体を初期化していることで時間がかかっていると思われます。以下のコードを試しました。
- t = 1 の場合、配列を初期化するのでなく、差分情報（配列 diff）だけ初期化する。
- t = 2 の場合、差分情報を記録する。
- t = 3 の場合、ベースの値と、差分情報から A[i] を出力する。
以下は、コードの抜粋です。このコードも、問題で示されている入力例に対して、正しい出力をします。
bool t1_occur = false;
ull base = 0;
ull diff[n + 1];
memset(diff, 0, sizeof(diff));
int q;
cin >> q;
for (int i = 1; i <= q; ++i) {
int t, iq, xq;
cin >> t;
if (t == 1) {
cin >> xq;
t1_occur = true;
base = xq;
memset(diff, 0, sizeof(diff));
} else if (t == 2) {
cin >> iq >> xq;
diff[iq] += xq;
} else {
cin >> iq;
if (t1_occur) {
cout << base + diff[iq] << endl;
} else {
cout << a[iq] + diff[iq] << endl;
}
}
}ただし、このプログラムも TLE と判定されました。初期値の代わりに差分情報の配列を初期化する必要があるため、計算量を下げることができていません。
バージョンと差分を管理する。
時間がかかるのは、結局、何回も配列を初期化しているためだと思われます。差分情報を分けて考えるのは良いアイデアで、ｔ = 1 を受け取った回数をバージョンとして合わせて管理します。
具体的には、以下となります。
- t = 1 を受け取った回数を g（gererations: 世代）とします。
- t = 1 の場合、g をインクリメントして、ベース値（base）を更新します。
- t = 2 の場合、g と i を組にして、差分を記録します。
- t = 3 の場合、g と i から差分を取り出し、ベース値に加えて出力します。
差分情報は、pair をキーとする map としました。この工夫で配列の初期化を無くすことができました。コードは以下となります。
#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long int ull;
int main()
{
int n;
cin >> n;
vector<ull> a(n + 1);
for (int i = 1; i <= n; ++i) {
cin >> a[i];
}
int g = 0;
ull base = 0;
map<pair<int, int>, ull> diff;
int q;
cin >> q;
for (int i = 1; i <= q; ++i) {
int t, iq, xq;
cin >> t;
if (t == 1) {
cin >> xq;
++g;
base = xq;
} else if (t == 2) {
cin >> iq >> xq;
diff[make_pair(g, iq)] += xq;
} else {
cin >> iq;
if (g > 0) {
cout << base + diff[make_pair(g, iq)] << endl;
} else {
cout << a[iq] + diff[make_pair(g, iq)] << endl;
}
}
}
return 0;
}このコードは、無事 AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC278D）
Python では、dict の代わりに defaultdict を使います。dict は、存在しないキーを指定したときエラーがでるためです。defaultdict は、C++ の map と同じように存在しないキーを指定してもエラーとならず、0 を返します。
基本的な考えは、C++ の map 版と同じです。
"""AtCoder Beginner Contest 278 D"""
from collections import defaultdict
n = int(input())
a = list(map(int, input().split()))
g = 0
base = 0
diff = defaultdict(int)
q = int(input())
for i in range(q):
t = list(map(int, input().split()))
if t[0] == 1:
g += 1
base = t[1]
elif t[0] == 2:
diff[(g, t[1])] += t[2]
else:
if g > 0:
print(base + diff[(g, t[1])])
else:
print(a[t[1] - 1] + diff[(g, t[1])])こちらも「AC」と判定されました。
最後に
この問題は、特定のアルゴリズムを適用するというより、配列やマップなどのコンテナの仕組みを効率よく使うことで解くことができました。
ABC278 は、A問題からD問題までは、実装力が問われる問題の傾向が強いと感じました。ある程度の数の問題数を解くことで、実装力を上げることができると考えています。
引き続き ABC の問題を紹介していきます。
