四平方定理

• 四平方定理

コンテスト中の考察

• 愚直だと$O(N^4)$で死
• x, y, zを決め打てば$O(N^3)$になる。Nの最大値は2000で、$N^3 = 8\times10^3$。ワンチャン通らねえかなぁと思って投げたらTLE。これに近い解法もTLE。この解法は無理っぽい
• 二分探索とかは無理そう。できても3つ決め打たないと無理っぽいので計算量が死。
• Prime Square Sumと似た雰囲気を感じるので考えてみる。
  • あの問題は$A^2 + B^2 + C^2 = N$の式を考える問題。Aに注目して考えると、$B^2 > 0, C^2 > 0$なので$A^2 < N$となる。よって$A < \sqrt{N}$の範囲で数字を列挙することで解けた。
  • この問題は$x^2 + y^2 + z^2 = w^2 + D$という式。$x$に注目すると$x^2 < w^2 + D$みたいになる。ただ、これをやっても計算量が減らない。なのでこのアプローチは無理。
• 式をこねくり回しても何も出てこない。もう一つ式を作ればなんかうまくいくかと思ったけど全然無理。えーん
• ようつべたのちい
• おしまい

解けなかった原因

• 値を一気に確定させないといけないと思ってしまっていた
  • x, y, zを一度に確定させないとwを確定させられずに死ぬって思ってた
  • x, yとz, wを別々に決めて合成するっていう発想がなかった。
  • 解法は半分全列挙らしいんだけど、「半分全列挙=bitをつかう」だと思ってた。アホ

解法

• 与式$x^2 + y^2 + z^2 = w^2 + D$を$x^2 + y^2 = w^2 - z^2 + D$に変形する
• すると、左辺に$x, y$で右辺に$w, z​$といった感じに変数が分割される。
• 左辺と右辺の組み合わせを別々に計算しておく。そんで、最後に左辺の組み合わせの数と右辺の組み合わせの数の積を多脚合わせていく。
• 計算量は$O(N^2)$となる

コード

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

#define int long long

int N, D;
int XY[2*2000*2000+100]; // XY[i]: 左辺がiになる組み合わせの数
int WZ[2*2000*2000+100]; // WZ[i]: 右辺がiになる組み合わせの数

signed main() {
  cin >> N >> D;
  // 左辺 (x^2 + y^2) の組み合わせを数える
  for (int x = 1; x <= N; x++) {
    for (int y = 1; y <= N; y++) {
      int left = x*x + y*y;
      XY[left]++;
    }
  }

  // 右辺 (w^2 - z^2 + D) の組み合わせを数える
  for (int w = 1; w <= N; w++) {
    for (int z = 1; z <= N; z++) {
      int right = w*w - z*z + D;
      // 式が複雑なので判定を入れておく
      if (right >= 2 and right <= 2*N*N) {
        WZ[right]++;
      }
    }
  }

  int ans = 0;
  // (値がiとなる左辺の場合の数 * 値がiとなる右辺の場合の数)を足し合わせる
  for (int i = 2; i <= 2*N*N; i++) {
    ans += XY[i] * WZ[i];
  }

  cout << ans << endl;

  return 0;
}

別解

• 本質部分はさっきと同じで式変形して半分全列挙する。片方は前計算で求めておいて、もう片方は二分探索で探す。

• 計算量は$O(N^2logN)$

• 4 Values whose Sum is 0と同じ解き方（蟻本に載ってる）

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

#define int long long

int N, D;
int AB[2222*2222];

signed main() {
  cin >> N >> D;

  int cnt = 0;
  for (int i = 1; i <= N; i++) {
    for (int j = 1; j <= N; j++) {
      AB[cnt] = i*i + j*j;
      cnt++;
    }
  }

  sort(AB, AB+cnt);

  int ans = 0;
  for (int i = 1; i <= N; i++) {
    for (int j = 1; j <= N; j++) {
      int num = i*i - j*j + D;
      int left = lower_bound(AB, AB+cnt, num) - AB;
      int right = upper_bound(AB, AB+cnt, num) - AB;
      ans += right - left;
    }
  }

  cout << ans << endl;

  return 0;
}

要点

• 等式の半分全列挙
  • 右辺と左辺で別々に値を計算して最後に合成する
  • 全列挙無理→3個決め打ち無理→半分全列挙考えるって感じだろうか。よくわからん

感想とか

• 難しい。解けなかった。えーん
• 100人近く通してるので日本は安泰。有能すぎる
• 「半分全列挙＝bitをつかう」だと思ってた。全列挙の半分が高速だったら半分全列挙を考える価値がありそう
• 蟻本の「くじびき」が類題。P8が問題説明でP25が解説。
• 蟻本P147の4 Values whose Sum is 0もこの問題の類題。