けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AISing Programming Contest 2020 C - XYZ Triplets (灰色, 300 点)

 n = 1, 2, \dots, N ごとに考えるのではなく、集計する考え方をするのがポイント!!!

問題概要

正の整数  N が与えられる。各  n = 1, 2, \dots, N に対して、

  •  x, y, z は正の整数
  •  x^{2} + y^{2} + z^{2} + xy + yz + zx = n

を満たすような  (x, y, z) の組の個数を求めよ。

制約

  •  1 \le N \le 10^{4}

考えたこと

最初に特定の整数  n に対して

  •  x, y, z は正の整数
  •  x^{2} + y^{2} + z^{2} + xy + yz + zx = n

を満たす  (x, y, z) の個数を求めることを考えてみる。まず、

 x^{2} \lt x^{2} + y^{2} + z^{2} + xy + yz + zx = n

より、 x \lt \sqrt{n} となることがわかる。同様に  y \lt \sqrt{n},  z \lt \sqrt{n} も成立している。よって、

  •  x = 1, 2, \dots, \sqrt{n}
  •  y = 1, 2, \dots, \sqrt{n}
  •  z = 1, 2, \dots, \sqrt{n}

をすべて試せば OK。この場合の計算量は  O(n\sqrt{n}) となる。

では、 n = 1, 2, \dots, N のすべてについて求めるにはどうしたらよいだろうか?一見するとまともに上記のことをやると  O(N^{2} \log{N}) かかりそうな気がしてくる。

しかしよく考えると、

  •  x = 1, 2, \dots, \sqrt{n}
  •  y = 1, 2, \dots, \sqrt{n}
  •  z = 1, 2, \dots, \sqrt{n}

をすべて試すときに、 v = x^{2} \lt x^{2} + y^{2} + z^{2} + xy + yz + zx として

counter[v]++;

という感じにすれば OK。こうすることで、 n = 1, 2, \dots, N に対する個数をまとめて求めることができる。よって全体の計算量は  O(N \sqrt{N}) でできる。

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

int main() {
    int N;
    cin >> N;
    vector<long long> num(N+1, 0);
    for (long long x = 1; x <= 100; ++x) {
        for (long long y = 1; y <= 100; ++y) {
            for (long long z = 1; z <= 100; ++z) {
                long long n = x*x + y*y + z*z + y*z + z*x + x*y;
                if (n <= N) num[n]++;
            }
        }
    }
    for (int n = 1; n <= N; ++n) cout << num[n] << endl;
}