正方形計數 題解

XuYueming發表於2024-09-22

題意簡述

給出一個 \(n \times n\) 的格點平面，有 \(q\) 次詢問，求有多少正方形以 \((x, y)\) 為某一頂點，滿足這個正方形頂點均在格點上，且邊長為有理數。

\(l \leq 10^5\)，\(q \leq 5 \times 10^5\)。

題目分析

看到邊長為有理數，想到「畢達哥拉斯三元組」（"Pythagorean triple"，簡稱「三元組」）。在 \(n\) 範圍內，這樣的三元組是不多的，根據下文的生成方法可以證明組數是 \(\Theta(n \log n)\) 的。我們先來討論如何生成出這些「三元組」。

本文中只介紹最簡單的「歐幾里得公式」（"Euclid's formula"）來生成所有 \(n\) 範圍內的「三元組」。如果讀者有興趣，可以嘗試閱讀 Formulas for generating Pythagorean triples，並使用其他方法來解決此問題。

設生成的「三元組」為 \((a, b, c)\)，那麼我們根據引數 \(\varphi, \lambda, k\)，其中 $\varphi \gt \lambda \gt 0, \lambda \perp \varphi , k \in\mathbb{Z}_+ $ 且 \(\lambda\) 和 $\varphi $ 中恰有一個偶數，根據下述公式即可生成所有「三元組」：

\({\displaystyle a=k\cdot (\varphi ^{2}-\lambda^{2}),\ \,b=k\cdot (2\varphi \lambda),\ \,c=k\cdot (\varphi ^{2}+\lambda^{2})}\) [1]

當 \(k = 1\) 時生成出來的 \((a, b, c)\) 為「本原三元組」（"primitive triple"），如 \({\displaystyle \varphi =2,\ \,\lambda=1,\, k = 1}\) 生成出 \((3, 4, 5)\)。所有「三元組」都能唯一地被某一「本原三元組」和縮放係數 \(k\) 生成。

使用初等代數展開，就能證明生成的演算法的充分性，這部分不在贅述。考慮證明其必要性，即所有「本原三元組」都能用這樣的方式表示出來。

證明 [1:1]：

「本原三元組」\((a, b, c)\) 滿足 \(a^2 + b^2 = c^2\)，且 \(\gcd(a, b, c) = 1\)。那麼有 \(a\)、\(b\)、\(c\) 兩兩互質，那麼 \(a\) 和 \(b\) 間至少一個是奇數，不妨令 \(a\) 為奇數，那麼 \(b\) 為偶數且 \(c\) 為奇數（如果 \(a\) 和 \(b\) 都是奇數，那麼 \(c^2\) 會是偶數，那麼 \(c\) 也會是偶數，那麼 \(4 \mid c^2\)，而 \(a^2 + b^2 \bmod 4 = 2\)，矛盾）。

我們得到 \(c^2 - a^2 = b^2\)，因此 \((c-a)(c+a) = b^2\)，即 \(\dfrac{c+a}{b} = \dfrac{b}{c-a}\)。\(\dfrac{c+a}{b}\) 顯然為有理數，設其最簡形式為 \(\dfrac{\varphi }{\lambda}\)。因此 \(\dfrac{c-a}{b} = \dfrac{\lambda}{\varphi }\)。我們有：

\[\frac{c}{b} + \frac{a}{b} = \frac{\varphi }{\lambda}, \quad \quad \frac{c}{b} - \frac{a}{b} = \frac{\lambda}{\varphi } \]

解出 \(\dfrac{c}{b}\) 和 \(\dfrac{a}{b}\) 分別為：

\[\frac{c}{b} = \frac{1}{2}\left(\frac{\varphi }{\lambda} + \frac{\lambda}{\varphi }\right) = \frac{\varphi ^2 + \lambda^2}{2\varphi \lambda}, \quad \quad \frac{a}{b} = \frac{1}{2}\left(\frac{\varphi }{\lambda} - \frac{\lambda}{\varphi }\right) = \frac{\varphi ^2 - \lambda^2}{2\varphi \lambda}. \]

由於 \(\varphi \perp \lambda\)，故它們不能同為偶數。若它們同為奇數，則 \(\dfrac{\varphi ^2 - \lambda^2}{2\varphi \lambda}\) 的分子將是 \(4\) 的倍數，而分母關於 \(2\) 的因子有且僅有一個 \(2\)，匯出 \(a\) 必為偶數，和我們假設 \(a\) 為奇數矛盾。因此，\(\lambda\) 和 $\varphi $ 中恰有一個偶數，所以 \(\varphi ^2 \pm \lambda^2\) 為奇數。因此，這些分數是最簡分數（如果有一個奇質數 \(p\) 為分子分母的公因數，由於 \(\lambda\) 和 $\varphi $ 的奇偶性，只能整除 \(\lambda\) 或 $\varphi $ 其中一個，那麼就不可以整除分子 \(\varphi ^2 \pm \lambda^2\)，矛盾）。由此將分子與分子相等，分母與分母相等，得到歐幾里得公式：

\[a = \varphi ^2 - \lambda^2, \quad b = 2\varphi \lambda, \quad c = \varphi ^2 + \lambda^2 \]

我們可以 \(\Theta(\sqrt{n})\) 地列舉 \(\varphi\)，再 \(\mathcal{O}(\sqrt{n})\) 的列舉 \(\lambda \lt \varphi\)，這也說明了 \(a, b, c \leq n\) 的「本原三元組」是 \(\Theta(\sqrt{n})\cdot\mathcal{O}(\sqrt{n}) = \mathcal{O}(n)\) 的。事實上，這是一個較松的上界。

根據 \(\varphi ^2 + \lambda^2 \leq c = k \cdot(\varphi ^2 + \lambda^2) \leq n\)，得到 \(\lambda \leq \sqrt{n - \varphi ^2}\)。同時結合 \(\lambda \lt \varphi\)。得到列舉 \(\lambda\) 更精確的時間複雜度是 \(\Theta \left(\min \left\{ \sqrt{n - \varphi ^2}, \varphi - 1 \right\}\right)\)。

由於 \(\min \left\{ \sqrt{n - \varphi ^2}, \varphi-1 \right\} \leq \varphi-1\)，我們不妨按照 \(\varphi - 1\) 計算上界。

此時內層為 \({\displaystyle \mathcal{O}\left(\sum_{\lambda=1}^{\varphi - 1} \dfrac{n}{\varphi ^2 + \lambda^2}\right) \leq \mathcal{O}\left(\sum_{\lambda=1}^{\varphi - 1} \dfrac{n}{\varphi ^2}\right) = \mathcal{O}\left(\dfrac{n}{\varphi}\right)}\)。其中第二步放縮是在分母中，直接去掉了 \(\lambda^2\) 這一項。我之前推的時候因為 \(\lambda \lt \varphi\)，把分母看做 \(2\lambda^2\)，然後省略常數什麼的，得到的上界是很鬆的 \(\Theta(n\sqrt{n})\)，疑惑了好久。

於是就能得到總的時間複雜度的式子最後長這樣：

\[\begin{aligned} &\quad\ \mathcal{O}\left(\sum_{\varphi=1}^{\sqrt{n}} \sum_{\lambda=1}^{\min \left\{ \sqrt{n - \varphi ^2}, \varphi - 1 \right\}} \dfrac{n}{\varphi ^2 + \lambda^2} \right) \\ & \leq \mathcal{O}\left(\sum_{\varphi=1}^{\sqrt{n}} \sum_{\lambda=1}^{\varphi-1} \dfrac{n}{\varphi ^2 + \lambda^2} \right) \\ & \leq \mathcal{O}\left(\sum_{\varphi=1}^{\sqrt{n}} \frac{n}{\varphi} \right) \\ & \leq \mathcal{O}\left(n\sum_{\varphi=1}^{n} \frac{1}{\varphi} \right) \\ & \leq \mathcal{O}\left(n \log n \right) \\ \end{aligned} \]

我們使用 \(\Theta(n \log n)\) 的演算法找出了 \(n\) 以內所有「三元組」。話說回來，生成出來有什麼用呢？我們發現一個斜著的正方形，可以透過「改斜歸正」，四條邊分別往外做四個全等的直角三角形，且三角形三邊為我們求出的「三元組」，變成一個正的正方形。為了方便討論，我們可以特殊地令「三元組」\((a, b, c)\) 其中可以有 \(0\)，也就是本來就是正的正方形，也可以這樣操作。

改斜歸正

注意到，對於 \(a, b \neq 0\)，交換 \(a, b\)，也就是將上圖左右翻轉，也是合法的。我們現在考慮 \((a, b, c)\) 對整個平面哪些點有貢獻。我們不妨欽定詢問的點，作為這些正方形的某一頂點，然後考慮邊界情況，可以得到下圖。

對某一矩形的貢獻

發現，如果詢問的點在紅色矩形內，就可以用這組 \((a, b, c)\) 做出以 \((x, y)\) 為左頂點的合法正方形。可以離線下來做掃描線。其它 \(3\) 個端點類似處理即可。總共 \(4\) 次掃描線。或者由於圖形很強的對稱性，不用做 \(4\) 次掃描線，只做 \(1\) 次來最佳化常數，讀者自行最佳化不難。

時間複雜度 \(\Theta(n \log^2n + q \log n)\)。

程式碼

#include <cstdio>
#include <iostream>
#include <vector>
#include <algorithm>
#include <cstring>
using namespace std;

const int N = 100010, Q = 500010;

int n, q, ans[Q];
vector<pair<int, int>> line[N], qx[N], qy[N], triple;

struct Bit_Tree {
    static constexpr inline int lowbit(int x) {
        return x & -x;
    }
    int tree[N];
    inline void clear() {
        memset(tree + 1, 0x00, sizeof (int) * (n));
    }
    inline void modify(int p, int v) {
        for (int i = p; i <= n; i += lowbit(i))
            tree[i] += v;
    }
    inline void modify(int l, int r, int v) {
        modify(l, v), modify(r + 1, -v);
    }
    inline int query(int p) {
        int res = 0;
        for (int i = p; i; i -= lowbit(i))
            res += tree[i];
        return res;
    }
} yzh;

signed main() {
    #ifndef XuYueming
    freopen("WSDR.in", "r", stdin);
    freopen("WSDR.out", "w", stdout);
    #endif
    scanf("%*d%d%d", &n, &q);
    for (int i = 1; i * i <= n; ++i)
    for (int j = (i & 1) ^ 1; i * i + j * j <= n && j < i; j += 2)
        if (__gcd(i, j) == 1) {
            for (int k = 1; k * (i * i + j * j) <= n; ++k) {
                int a = k * (i * i - j * j), b = k * (2 * i * j);
                if (a + b + 1 <= n) {
                    triple.emplace_back(a, b);
                    if (b) triple.emplace_back(b, a);
                    // 交換 (a, b)
                }
            }
        }
    for (int i = 1, x, y; i <= q; ++i) {
        scanf("%d%d", &x, &y);
        qx[x].emplace_back(y, i);
        qy[y].emplace_back(x, i);
    }
    
    for (auto& [a, b]: triple) line[n - a - b].emplace_back(1 + a, n - b);
    for (int i = n; i >= 1; --i) {
        for (auto& [l, r]: line[i]) yzh.modify(l, r, 1);
        for (auto& [y, p]: qx[i]) ans[p] += yzh.query(y);
    }
    
    yzh.clear();
    for (int i = 1; i <= n; ++i) line[i].clear();
    for (auto& [a, b]: triple) line[1 + a + b].emplace_back(1 + b, n - a);
    for (int i = 1; i <= n; ++i) {
        for (auto& [l, r]: line[i]) yzh.modify(l, r, 1);
        for (auto& [y, p]: qx[i]) ans[p] += yzh.query(y);
    }
    
    yzh.clear();
    for (int i = 1; i <= n; ++i) line[i].clear();
    for (auto& [a, b]: triple) line[1 + a + b].emplace_back(1 + a, n - b);
    for (int i = 1; i <= n; ++i) {
        for (auto& [l, r]: line[i]) yzh.modify(l, r, 1);
        for (auto& [y, p]: qy[i]) ans[p] += yzh.query(y);
    }
    
    yzh.clear();
    for (int i = 1; i <= n; ++i) line[i].clear();
    for (auto& [a, b]: triple) line[n - a - b].emplace_back(1 + b, n - a);
    for (int i = n; i >= 1; --i) {
        for (auto& [l, r]: line[i]) yzh.modify(l, r, 1);
        for (auto& [y, p]: qy[i]) ans[p] += yzh.query(y);
    }
    
    for (int i = 1; i <= q; ++i)
        printf("%d\n", ans[i]);
    return 0;
}

References

  1. Wikipedia, "Pythagorean Triple", Generating a Triple and Proof of Euclid's Formula。 ↩︎ ↩︎

相關文章