I wrote some code to find positive integer solutions (x, y) to the Pell equation x^2 - d\cdot y^2 = 1 for a user-inputted value d. Let the smallest (trivial) solution be (a, b). All solutions x + y\sqrt{d} to the Pell equation must be of the form \left(a+b\sqrt{d}\right)^n. In my code, I first brute force the trivial solution (a, b) and then use Dynamic Programming to iteratively craft the next solution by multiplying (x+y\sqrt{d}) \cdot (a+b\sqrt{d}). To find square roots for the brute force part, I use binary search to find the floor of \sqrt{X}.
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define pell cout << "Integer solutions to x^2 - dy^2 = 1:" << endl
void perfect_square(){
cout << "d is a perfect square" << endl;
pell;
cout << "1 0" << endl;
exit(0);
}
// returns the floor of sqrt(x)
ll sqrt(ll x, ll p, ll q){
if(p >= q) return p;
ll mid = p + (q-p+1)/2;
if(mid*mid <= x) return sqrt(x, mid, q);
else return sqrt(x, p, mid-1);
}
int main() {
int d; ll a, b, x1, x2, y1, y2;
cin >> d; // find solutions to x^2 - dy^2 = 1
if(sqrt(d,0,d)*sqrt(d,0,d)==d)
perfect_square();
pell;
vector<vector<ll>> dp = {{1,0}};
for(b=1;true;b++){
ll a2 = 1+d*b*b;
a = sqrt(a2, 0, a2);
if(a*a == a2){
// (x+y[d])(a+b[d]) = (a*x+b*d*y) + (b*x+a*y)[d]
x1 = a, x2 = b*d, y1 = b, y2 = a;
break;
}
}
for(ll i=1;true;i++){
ll x = dp[i-1][0], y = dp[i-1][1]; // previous solution
cout << x << " " << y << endl;
ll l = LLONG_MAX;
if(x > l/max(a,b)) // are x1*x or y1*x out of bounds
break;
if(y > l/max(b*d,a)) // are x2*y or y2*y out of bounds
break;
if(x1*x + x2*y < x1*x) // is x1*x + x2*y out of bounds
break;
if(y1*x + y2*y < y1*x) // is y1*x + y2*y out of bounds
break;
dp.push_back({x1*x + x2*y, y1*x + y2*y});
}
}
Please feel free and encouraged to provide any feedback on my above code.
To check if the next solution \{x_1\cdot x + x_2\cdot y, y_1\cdot x + y_2\cdot y\} is within the bounds of long long, I had to do some tedious casework. Can we check if x_1\cdot x + x_2\cdot y is out-of-bounds more efficiently? Do we really need to use 4 separate if-else statements to check out-of-bounds, or can you please provide an easier method?
I am also uncertain as to whether binary searching on the floor of the square root of X is really necessary. I only use binary search to find the square root of a number if it is a perfect square. Can you please explain how to find the square root faster using Silver-level algorithms? If you choose to reference a pre-existing C++ function for square root, please explain the algorithm used rather than referring me directly to that in-built function.
Thanks!