藍橋杯-四平方和

四平方和定理，又稱為拉格朗日定理：

每個正整數都可以表示為至多 4 個正整數的平方和。

如果把 0 包括進去，就正好可以表示為 4 個數的平方和。

比如：

5= 0^2 + 0^2 + 1^2 + 2^2

7= 1^2 + 1^2 + 1^2 + 2^2

對於一個給定的正整數，可能存在多種平方和的表示法。

要求你對 4 個數排序：

0≤a≤b≤c≤d

並對所有的可能表示法按 a,b,c,d 為聯合主鍵升序排列，最後輸出第一個表示法。

輸入格式

輸入一個正整數 N。

輸出格式

輸出4個非負整數，按從小到大排序，中間用空格分開。

資料範圍

0<N<5∗1e6

輸入樣例：

5

輸出樣例：

0 0 1 2

題解:

~~暴力three重迴圈判斷會超時~~~

列舉 c 和 d 的所有可能, 並把 c, d, c * c + d * d 都存到陣列中
對陣列排序。 (從大到小, 從小到達都行, 但對應的二分不一樣)
列舉 a 和 b 的所有可能, 用二分找當前 a b 下, 對應的 c d, 判斷是否滿足 n = a * a + b * b + c * c + d * d, 滿足的話輸出並退出ok了

#include <bits/stdc++.h>
using namespace std;
const int N = 1e7 + 10;
struct Node {
    int c, d, sum;
};
Node v[N];
bool cmp(Node a, Node b)
{
    if (a.sum != b.sum) return a.sum < b.sum;
    if (a.c != b.c) return a.c < b.c;
    return a.d < b.d;
}

int main()
{
    int n, m = 0; cin >> n;
    for (int i = 0; i * i <= n; i ++)
        for (int j = i; i * i + j * j <= n; j ++)
            v[m ++] = {i, j, i * i + j * j};
    sort(v, v + m, cmp);  // 從小到大排序
    
    for (int i = 0; i * i <= n; i ++)
        for (int j = 0; j * j + i * i <= n; j ++)  // a b 要從 0 開始列舉, 題中要求a≤b≤c≤d
        {
            int l = 0, r = m - 1;
            int t = n - i * i - j * j;
            while (l < r)
            {
                int mid = l + r >> 1;
                if (v[mid].sum >= t) r = mid;  // 二分, 相同的v[mid].sum, 查詢的是第一個滿足的, 由於我們是 升序, 剛好滿足 c <= d
                else l = mid + 1;
            }
            if (v[l].sum == t)
            {
                cout << i << ' ' << j << ' ' << v[l].c << ' ' << v[l].d << endl;
                return 0;
            }
        }
    return 0;
}

~~覺得寫的不錯的話, 點個贊吧~~~

下面不是題解, 是筆者為了讓自己更好理解二分寫的筆記

對比上面的程式碼可以發現程式碼中有 "<--------" 註釋的是發生改變的
其中 if (v[mid].sum <= t), 二分條件發生變化, 導致了這個二分在 v[mid].sum 相同的情況下, 找到的是 最後一個 符合條件的下標
上面的程式碼在 v[mid].sum 相同的情況下, 找到的是 第一個 符合條件的下標
so 在下面程式碼排序中我們是讓 c 和 d 按照倒序的方法排的順序

#include <bits/stdc++.h>
using namespace std;
const int N = 1e7 + 10;
struct Node {
    int c, d, sum;
};
Node v[N];
bool cmp(Node a, Node b)
{
    if (a.sum != b.sum) return a.sum < b.sum;
    if (a.c != b.c) return a.c > b.c;  // <------
    return a.d > b.d;       // <-------
}

int main()
{
    int n, m = 0; cin >> n;
    for (int i = 0; i * i <= n; i ++)
        for (int j = i; i * i + j * j <= n; j ++)
            v[m ++] = {i, j, i * i + j * j};
    sort(v, v + m, cmp);
    
    for (int i = 0; i * i <= n; i ++)
        for (int j = 0; j * j + i * i <= n; j ++)
        {
            int l = 0, r = m - 1;
            int t = n - i * i - j * j;
            while (l < r)
            {
                int mid = l + r + 1 >> 1;
                if (v[mid].sum <= t) l = mid;   // <-------
                else r = mid - 1;   // <-------
            }
            if (v[l].sum == t)
            {
                cout << i << ' ' << j << ' ' << v[l].c << ' ' << v[l].d << endl;
                return 0;
            }
        }
    return 0;
}

口訣
if 後面是 r 的話找到的是 l 的, 即是 第一個 滿足條件的
if 後面是 l 的話找到的是 r 的, 即是 最後一個 滿足條件的

藍橋杯-四平方和

相關文章