$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;
}

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

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

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

問題概要

制約

考えたこと