AtCoder が提供しているABC(AtCoder Beginner Contest)366 D問題をC++とPythonで解いてみました。ABC366は、2024年8月10日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 Cuboid Sum Query(Difficulty : 586)
問題の詳細は、リンク先をご覧ください。
三次元の累積和を使う典型問題です。AtCoder Problems による Difficulty は 586 でした。
解答案
2次元の累積和の定番コードは、以下です。
int h, w;
cin >> h >> w;
vector<vector<int>> x(h, vector<int>(w));
for (int i = 0; i < h; ++i) {
for (int j = 0; j < w; ++j) {
cin >> x[i][j];
}
}
for (int i = 0; i < h; ++i) {
for (int j = 0; j < w; ++j) {
s[i + 1][j + 1] = s[i + 1][j] + x[i][j];
}
}
for (int i = 0; i < h; ++i) {
for (int j = 0; j < w; ++j) {
s[i + 1][j + 1] += s[i][j + 1];
}
}
// [a, c)*[b, d) の合計を求める。
cin >> a >> b >> c >> d;
result = s[c][d] - s[a][d] - s[c][b] + s[a][b];
3次元の場合、次元がひとつ増えますが、考え方は同じです。
C++ プログラム例(ABC366D)
プログラムの補足をします。
- 3次元配列 a を読み込みます(10-17行目)。
- 累積和 s を求めます。配列の要素が一つ増えるのは2次元の場合と同じです(19ー40行目)。
- クエリを読み込み、累積和から [Lx, Rx) × [Ly, Ry) × [Lz, Rz) の値を求めます。クエリで与えられる6種類の数字は、1から始まり、累積和の計算は0から始めます。このため、Lx/Ly/Lz を一つ引いて計算します。Rの値は一つ減らして右を開区間にするため一つ加えるため、値をそのまま使います(47ー49行目)。
- 値を求める式が長くなるため分けて書きました(50ー58行目)。
以下が、C++プログラムです。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
int main()
{
int n;
cin >> n;
vector<vector<vector<int>>> a(n, vector<vector<int>>(n, vector<int>(n, 0)));
for (int i = 0; i < n; ++i) {
for (int j = 0; j < n; ++j) {
for (int k = 0; k < n; ++k) {
cin >> a[i][j][k];
}
}
}
vector<vector<vector<int>>> s(n + 1, vector<vector<int>>(n + 1, vector<int>(n + 1, 0)));
for (int i = 0; i < n; ++i) {
for (int j = 0; j < n; ++j) {
for (int k = 0; k < n; ++k) {
s[i + 1][j + 1][k + 1] = s[i + 1][j + 1][k] + a[i][j][k];
}
}
}
for (int i = 0; i < n; ++i) {
for (int j = 0; j < n; ++j) {
for (int k = 0; k < n; ++k) {
s[i + 1][j + 1][k + 1] += s[i + 1][j][k + 1];
}
}
}
for (int i = 0; i < n; ++i) {
for (int j = 0; j < n; ++j) {
for (int k = 0; k < n; ++k) {
s[i + 1][j + 1][k + 1] += s[i][j + 1][k + 1];
}
}
}
int q;
cin >> q;
for (int i = 0; i < q; ++i) {
int Lx, Rx, Ly, Ry, Lz, Rz;
cin >> Lx >> Rx >> Ly >> Ry >> Lz >> Rz;
--Lx;
--Ly;
--Lz;
ll result = 0;
result += s[Rx][Ry][Rz];
result -= s[Lx][Ry][Rz];
result -= s[Rx][Ly][Rz];
result -= s[Rx][Ry][Lz];
result += s[Lx][Ly][Rz];
result += s[Lx][Ry][Lz];
result += s[Rx][Ly][Lz];
result -= s[Lx][Ly][Lz];
cout << result << endl;
}
return 0;
}
AC(Accepted=正しいプログラム)と判定されました。
Python プログラム例(ABC366D)
Python版も基本的な考え方は同じです。以下がプログラムです。
"""AtCoder Beginner Contest 366 D"""
n = int(input())
a = [[[] for _ in range(n)] for _ in range(n)]
for i in range(n):
for j in range(n):
a[i][j] = list(map(int, input().split()))
s = [[[0 for _ in range(n + 1)] for _ in range(n + 1)] for _ in range(n + 1)]
for i in range(n):
for j in range(n):
for k in range(n):
s[i + 1][j + 1][k + 1] = s[i + 1][j + 1][k] + a[i][j][k]
for i in range(n):
for j in range(n):
for k in range(n):
s[i + 1][j + 1][k + 1] += s[i + 1][j][k + 1]
for i in range(n):
for j in range(n):
for k in range(n):
s[i + 1][j + 1][k + 1] += s[i][j + 1][k + 1]
q = int(input())
for i in range(q):
Lx, Rx, Ly, Ry, Lz, Rz = map(int, input().split())
Lx -= 1
Ly -= 1
Lz -= 1
result = 0
result += s[Rx][Ry][Rz]
result -= s[Lx][Ry][Rz]
result -= s[Rx][Ly][Rz]
result -= s[Rx][Ry][Lz]
result += s[Lx][Ly][Rz]
result += s[Lx][Ry][Lz]
result += s[Rx][Ly][Lz]
result -= s[Lx][Ly][Lz]
print(result)
こちらも「AC」と判定されました。
最後に
2次元の累積和を使う問題は、一工夫必要な場合が多いです。この問題は3次元になった代わりに累積和をそのまま使う問題になっています。
引き続き ABC の問題を紹介していきます。