AOJ 2678 Cube Coloring

日本語記事が見つからなかったので

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2678

Cube Coloring | Aizu Online Judge

1 × 1 × 1 の小立方体で構成される X × Y × Z の直方体と座標P (A, B, C) が与えられる.

各小立方体の座標を (0, 0, 0) から (X-1, Y-1, Z-1) までの座標で表し, 小立方体 (x1, y1, z1) と (x2, y2, z2) の距離を |x1 - x2| + |y1 - y2| + |z1 - z2| で定義する.

正整数 N が与えられるので 各 i+1 (1 <= i + 1 <= N) に対し, P との距離d % N が i となるような小立方体の個数を出力せよ.  (2 <= N <= 1000, 1 <= X, Y, Z <= 1e6)

解法

O( X * Y * Z ) で全ての小立方体を数えていては TLE する.

各座標軸で mod を取り複数の小立方体をまとめた後, 定義に従って数え上げればよい.

計算量は O(N3).

ソースコード

#include "bits/stdc++.h"

using namespace std;

using ll = long long;

ll X, Y, Z, A, B, C, N;
ll cnt_x[1000], cnt_y[1000], cnt_z[1000], res[1000];

int main()
{
    cin.tie(0);
    ios::sync_with_stdio(false);

    cin >> X >> Y >> Z >> A >> B >> C >> N;
    for (int i = 0; i < X; i++)
    {
        int diff = abs(i - A) % N;
        cnt_x[diff]++;
    }
    for (int i = 0; i < Y; i++)
    {
        int diff = abs(i - B) % N;
        cnt_y[diff]++;
    }
    for (int i = 0; i < Z; i++)
    {
        int diff = abs(i - C) % N;
        cnt_z[diff]++;
    }
    for (int i = 0; i < N; i++)
    {
        for (int j = 0; j < N; j++)
        {
            for (int k = 0; k < N; k++)
            {
                int idx = i + j + k;
                if (idx >= 2 * N)
                    idx -= 2 * N;
                else if (idx >= N)
                    idx -= N;
                res[idx] += cnt_x[i] * cnt_y[j] * cnt_z[k];
            }
        }
    }
    for (int i = 0; i < N; i++)
    {
        cout << res[i] << " \n"[i == N - 1];
    }
}

感想

3重ループ内で % 演算子を使ったら TLE したので if 文で処理したら4倍弱速くなって(3.9s -> 1.1s) AC. うーん.