kazuma8128’s blog

競プロの面白い問題を解きます

HackerRank Cube Summation

問題概要

N x N x N の3次元配列がある. (N <= 100)
一点更新クエリと, 範囲和クエリがくるので処理せよ.

解法

Nが小さいので 3 次元の Fenwick Tree を書けば終わりです.

3次元は流石にライブラリ化してなかったのでコンテストとかで出たらタイピングバトルになって面白そう(いいえ)

ソースコード

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

template <typename T>
class fenwick_tree3 {
    const int a, b, c;
    vector<vector<vector<T>>> data;
public:
    fenwick_tree3(int a_, int b_, int c_)
        : a(a_), b(b_), c(c_), data(a, vector<vector<T>>(b, vector<T>(c))) {}

    // [0, pi] & [0, pj] & [0, pk]
    T find(int pi, int pj, int pk) const {
        T res = 0;
        for (int i = pi; i >= 0; i = (i & (i + 1)) - 1) {
            for (int j = pj; j >= 0; j = (j & (j + 1)) - 1) {
                for (int k = pk; k >= 0; k = (k & (k + 1)) - 1) {
                    res += data[i][j][k];
                }
            }
        }
        return res;
    }

    // [li, ri] & [lj, rj] & [lk, rk]
    T find(int li, int lj, int lk, int ri, int rj, int rk) const {
        --li, --lj, --lk;
        return find(ri, rj, rk) - find(ri, rj, lk) - find(ri, lj, rk) - find(li, rj, rk)
             + find(ri, lj, lk) + find(li, rj, lk) + find(li, lj, rk) - find(li, lj, lk);
    }

    void add(int pi, int pj, int pk, T val) {
        for (int i = pi; i < a; i |= i + 1) {
            for (int j = pj; j < b; j |= j + 1) {
                for (int k = pk; k < c; k |= k + 1) {
                    data[i][j][k] += val;
                }
            }
        }
    }

    void update(int pi, int pj, int pk, T val) {
        add(pi, pj, pk, val - find(pi, pj, pk, pi, pj, pk));
    }
};

int main()
{
    ios::sync_with_stdio(false), cin.tie(0);
    int T;
    cin >> T;
    while (T--) {
        int N, M;
        cin >> N >> M;
        fenwick_tree3<ll> ft(N, N, N);
        while (M--) {
            string com;
            cin >> com;
            if (com == "UPDATE") {
                int x, y, z, w;
                cin >> x >> y >> z >> w;
                x--, y--, z--;
                ft.update(x, y, z, w);
            }
            else {
                int x1, y1, z1, x2, y2, z2;
                cin >> x1 >> y1 >> z1 >> x2 >> y2 >> z2;
                x1--, y1--, z1--, x2--, y2--, z2--;
                cout << ft.find(x1, y1, z1, x2, y2, z2) << endl;
            }
        }
    }
    return 0;
}