AtCoder が提供しているABC（AtCoder Beginner Contest）271 のB問題をC++とPythonで解いてみました。ABC271は、2022年10月1日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
B問題 Maintain Multiple Sequences（Difficulty : 97）
問題はリンク先をご覧ください。
ABC271 B問題 Maintain Multiple Sequences
数列を要素とする配列に対する操作が必要になります。AtCoder Problems による Difficulty は、97 でした。
解答案
問題を解く方針を書きだします。
- 数列を読み込む
- クエリを読み込み、sk 番目の数列の tk 項を出力する。
C++ の場合、vector の知識があると楽に解けます。
C++ プログラム例（ABC271B）
うまくいかないプログラムを先に紹介します。数列全体を2次元配列として、取り込もうとしています。
#include <bits/stdc++.h>
using namespace std;
#define N_MAX 200001
int a[N_MAX][N_MAX];
int main()
{
int n, q;
cin >> n >> q;
for (int i = 0; i < n; ++i) {
int l;
cin >> l;
for (int j = 0; j < l; ++j) {
cin >> a[i][j];
}
}
for (int k = 0; k < q; ++k) {
int i, j;
cin >> i >> j;
cout << a[i-1][j-1] << endl;
}
return 0;
}数列全体を格納する2次元配列 a のサイズが大きすぎて CE（Compile Error）の判定となります。なお、このプログラムは、N_MAX の値が大きくなければ正しく動作します。
問題を読むと、数列の長さは数列によって異なります。数列の長さの合計に次の制約があります。
$$\sum^N_{i=1} L_i \leqq 2 \times 10^5$$
1次元の配列として格納することにします。それぞれの数列の最初のインデックスを配列 sum_L として格納しておきます。配列のインデックスは、0 から開始ですが、問題では 1 から開始として与えられるため注意が必要です。
#include <bits/stdc++.h>
using namespace std;
#define N_MAX 200001
int a[N_MAX];
int L[N_MAX];
int sum_L[N_MAX];
int main()
{
int N, Q;
cin >> N >> Q;
int index = 0;
sum_L[0] = 0;
for (int i = 0; i < N; ++i) {
cin >> L[i];
if (i >= 1) {
sum_L[i] = sum_L[i - 1] + L[i - 1];
}
for (int j = 0; j < L[i]; ++j) {
cin >> a[index];
++index;
}
}
for (int k = 0; k < Q; ++k) {
int i, j;
cin >> i >> j;
cout << a[sum_L[i-1] + (j-1)] << endl;
}
return 0;
}Vector の知識があると、以下のような自然な実装ができます。数列の大きさ L を読み込み、vector コンテナのサイズを resize しています。
#include <bits/stdc++.h>
using namespace std;
#define N_MAX 200001
int main()
{
int n, q;
cin >> n >> q;
vector <vector<int>> a(n);
for (int i = 0; i < n; ++i) {
int l;
cin >> l;
a[i].resize(l);
for (int j = 0; j < l; ++j) {
cin >> a[i][j];
}
}
for (int k = 0; k < q; ++k) {
int i, j;
cin >> i >> j;
cout << a[i-1][j-1] << endl;
}
return 0;
}どちらも AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC271B）
リストは、C++ の配列と異なり、サイズは決まっていません。リスト a に数列を格納していきます。読み込むときにスライス機能を使って、最初の要素Liを無視しています。
"""AtCoder Beginner Contest 271 B"""
n, q = map(int, input().split())
a = [input().split()[1:] for i in range(n)]
for i in range(q):
s, t = map(int, input().split())
print(a[s - 1][t - 1])Python 版も「AC」と判定されました。
最後に
C++ で解く場合、単純に2次元配列が取れる大きさではなかったです。1次元配列に変換して解く方法と、Vector の resize でサイズを変更しながら格納する方法を紹介しました。
1次元配列に変換する方法は、細かい配慮が必要になります。単純にコンテナの機能を使う方法が好ましいと考えています。
引き続き ABC の問題を紹介していきます。
