Math/Pure/pell_equation.h

• View this file on GitHub
• Last update: 2022-12-27 16:11:10+08:00
• Link: https://github.com/zimpha/algorithmic-library/blob/master/cpp/mathematics/pell.py

Depends on

• Misc/int128.h

Code

#include "../../Misc/int128.h"

// From https://github.com/zimpha/algorithmic-library/blob/master/cpp/mathematics/pell.py
// Find minimum solution for x^2 - d*y^2 = 1 (d <= 1000)
// Pell equation {{{
void up(i128& ai, i128& aim, i128 alpha) {
    i128 old_ai = ai;
    ai = alpha * ai + aim;
    aim = old_ai;
}
pair<vector<i128>,i128> pqa(i128 p0, i128 q0, int d) {
    double sqrt_d = sqrt(d);
    i128 p_i = p0, q_i = q0;
    i128 p_ir = LLONG_MIN, q_ir = LLONG_MIN;
    i128 i = -1, ir = LLONG_MIN;
    vector<i128> alphas;

    while (1) {
        ++i;
        double xi_i = (p_i + sqrt_d) / q_i;
        double xibar_i = (p_i - sqrt_d) / q_i;
        i128 alpha_i = xi_i + 1e-9;

        if (ir == LLONG_MIN && 1 < xi_i && -1 < xibar_i && xibar_i < 0) {
            ir = i;
            p_ir = p_i;
            q_ir = q_i;
        }
        if (ir != LLONG_MIN && ir != i && p_i == p_ir && q_i == q_ir) {
            break;
        }
        alphas.push_back(alpha_i);
        p_i = alpha_i * q_i - p_i;
        q_i = (d - p_i * p_i) / q_i;
    }
    return {alphas, i - ir};
}

// return minimum solution for x^2 - d*y^2 = 1
pair<i128,i128> pell(int d) {
    auto [alphas, l] = pqa(0, 1, d);

    int index = (l % 2 == 1) ? 2*l-1 : l-1;

    i128 b_i = 0, b_im = 1;
    i128 g_i = 1, g_im = 0;
    int pl = alphas.size() - l;

    for (int i = 0; i <= index; ++i) {
        i128 alpha_i;
        if (i < pl) alpha_i = alphas[i];
        else alpha_i = alphas[pl + (i - pl) % l];

        up(b_i, b_im, alpha_i);
        up(g_i, g_im, alpha_i);
    }
    return {g_i, b_i};
}
// }}}

#line 1 "Misc/int128.h"
// i128 helper functions {{{
using i128 = __int128_t;
i128 str2i128(std::string str) {
    i128 ret = 0;
    bool minus = false;
    for (auto c : str) {
        if (c == '-')
            minus = true;
        else
            ret = ret * 10 + c - '0';
    }
    return minus ? -ret : ret;
}
std::istream &operator>>(std::istream &is, i128 &x) {
    std::string s;
    return is >> s, x = str2i128(s), is;
}
std::ostream &operator<<(std::ostream &os, const i128 &x) {
    i128 tmp = x;
    if (tmp == 0) return os << 0;
    std::vector<int> ds;
    if (tmp < 0) {
        os << '-';
        while (tmp) {
            int d = tmp % 10;
            if (d > 0) d -= 10;
            ds.emplace_back(-d), tmp = (tmp - d) / 10;
        }
    } else {
        while (tmp) ds.emplace_back(tmp % 10), tmp /= 10;
    }
    std::reverse(ds.begin(), ds.end());
    for (auto i : ds) os << i;
    return os;
}
i128 my_abs(i128 n) {
    if (n < 0) return -n;
    return n;
}
i128 gcd(i128 a, i128 b) {
    if (b == 0) return a;
    return gcd(b, a % b);
}
// Count trailing zeroes
int ctz128(i128 n) {
    if (!n) return 128;
 
    if (!static_cast<uint64_t>(n)) {
        return __builtin_ctzll(static_cast<uint64_t>(n >> 64)) + 64;
    } else {
        return __builtin_ctzll(static_cast<uint64_t>(n));
    }
}
// }}}

#line 2 "Math/Pure/pell_equation.h"

// From https://github.com/zimpha/algorithmic-library/blob/master/cpp/mathematics/pell.py
// Find minimum solution for x^2 - d*y^2 = 1 (d <= 1000)
// Pell equation {{{
void up(i128& ai, i128& aim, i128 alpha) {
    i128 old_ai = ai;
    ai = alpha * ai + aim;
    aim = old_ai;
}
pair<vector<i128>,i128> pqa(i128 p0, i128 q0, int d) {
    double sqrt_d = sqrt(d);
    i128 p_i = p0, q_i = q0;
    i128 p_ir = LLONG_MIN, q_ir = LLONG_MIN;
    i128 i = -1, ir = LLONG_MIN;
    vector<i128> alphas;

    while (1) {
        ++i;
        double xi_i = (p_i + sqrt_d) / q_i;
        double xibar_i = (p_i - sqrt_d) / q_i;
        i128 alpha_i = xi_i + 1e-9;

        if (ir == LLONG_MIN && 1 < xi_i && -1 < xibar_i && xibar_i < 0) {
            ir = i;
            p_ir = p_i;
            q_ir = q_i;
        }
        if (ir != LLONG_MIN && ir != i && p_i == p_ir && q_i == q_ir) {
            break;
        }
        alphas.push_back(alpha_i);
        p_i = alpha_i * q_i - p_i;
        q_i = (d - p_i * p_i) / q_i;
    }
    return {alphas, i - ir};
}

// return minimum solution for x^2 - d*y^2 = 1
pair<i128,i128> pell(int d) {
    auto [alphas, l] = pqa(0, 1, d);

    int index = (l % 2 == 1) ? 2*l-1 : l-1;

    i128 b_i = 0, b_im = 1;
    i128 g_i = 1, g_im = 0;
    int pl = alphas.size() - l;

    for (int i = 0; i <= index; ++i) {
        i128 alpha_i;
        if (i < pl) alpha_i = alphas[i];
        else alpha_i = alphas[pl + (i - pl) % l];

        up(b_i, b_im, alpha_i);
        up(g_i, g_im, alpha_i);
    }
    return {g_i, b_i};
}
// }}}

Back to top page