Problem 833: Square Triangular Numbers

Mathematical Foundation

Theorem 1 (Reduction to Pell’s Equation). The equation $n(n+1)/2 = m^2$ holds if and only if $(x, y) = (2n+1, 2m)$ satisfies the Pell equation $x^2 - 2y^2 = 1$.

Proof. Starting from $n(n+1) = 2m^2$, multiply both sides by 4:

Setting $x = 2n+1$ and $y = 2m$ gives $x^2 - 2y^2 = 1$. Conversely, any solution $(x,y)$ to $x^2 - 2y^2 = 1$ with $x$ odd and $y$ even yields $n = (x-1)/2 \in \mathbb{Z}_{\ge 0}$ and $m = y/2 \in \mathbb{Z}_{\ge 0}$ with $T_n = m^2$. $\square$

Theorem 2 (Structure of Pell Solutions). The equation $x^2 - 2y^2 = 1$ has infinitely many positive solutions. The fundamental solution is $(x_1, y_1) = (3, 2)$, and all positive solutions are given by

Proof. Existence of the fundamental solution: one verifies $3^2 - 2 \cdot 2^2 = 9 - 8 = 1$. For completeness, $(3,2)$ is fundamental because it is found from the continued fraction expansion $\sqrt{2} = [1; \overline{2}]$, whose first convergent $p_1/q_1 = 3/2$ satisfies $p_1^2 - 2q_1^2 = 1$.

To show all solutions arise as stated, note that the norm map $N(x + y\sqrt{2}) = x^2 - 2y^2$ is multiplicative: $N(\alpha\beta) = N(\alpha)N(\beta)$. If $(x_k, y_k)$ is a solution, then $N(x_k + y_k\sqrt{2}) = 1$, and $(x_k + y_k\sqrt{2})(3 + 2\sqrt{2}) = x_{k+1} + y_{k+1}\sqrt{2}$ also has norm 1. Conversely, if $(x,y)$ is any solution with $x + y\sqrt{2} > 1$, then $x + y\sqrt{2} \ge 3 + 2\sqrt{2}$ (since $3 + 2\sqrt{2}$ is the smallest unit $> 1$ in $\mathbb{Z}[\sqrt{2}]$), and dividing by $3 + 2\sqrt{2}$ yields a smaller solution. By induction, every solution is a power of the fundamental one. $\square$

Lemma (Recurrence). The Pell solutions satisfy the linear recurrences:

with initial conditions $(x_0, y_0) = (1, 0)$, $(x_1, y_1) = (3, 2)$.

Proof. From $(x_k + y_k\sqrt{2}) = (3+2\sqrt{2})^k$, the minimal polynomial of $\alpha = 3 + 2\sqrt{2}$ is $t^2 - 6t + 1 = 0$ (since $\alpha + \bar{\alpha} = 6$ and $\alpha\bar{\alpha} = 1$). Therefore $\alpha^{k+1} = 6\alpha^k - \alpha^{k-1}$, and equating rational and irrational parts gives the stated recurrences. $\square$

Theorem 3 (Parity Preservation). For all $k \ge 1$, $x_k$ is odd and $y_k$ is even.

Proof. By induction. Base: $x_1 = 3$ (odd), $y_1 = 2$ (even). The coupled recurrence $x_{k+1} = 3x_k + 4y_k$, $y_{k+1} = 2x_k + 3y_k$ shows: if $x_k$ is odd and $y_k$ is even, then $x_{k+1} = 3(\text{odd}) + 4(\text{even}) = \text{odd}$ and $y_{k+1} = 2(\text{odd}) + 3(\text{even}) = \text{even}$. $\square$

Editorial

Alternatively, for computing a sum of the first $K$ square triangular numbers, iterate the recurrence and accumulate. We compute K-th square triangular number mod p. Finally, method: matrix exponentiation.

Pseudocode

Compute K-th square triangular number mod p
Method: matrix exponentiation

Complexity Analysis

• Time: $O(\log K)$ via $2 \times 2$ matrix exponentiation modulo $p$; or $O(K)$ via direct iteration.
• Space: $O(1)$.

Answer

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

/*
 * Problem 833: Square Triangular Numbers
 *
 * Pell equation x^2-2y^2=1
 * Answer: 541476798
 */

const long long MOD = 1e9 + 7;

long long power(long long base, long long exp, long long mod) {
    long long result = 1;
    base %= mod;
    while (exp > 0) {
        if (exp & 1) result = result * base % mod;
        base = base * base % mod;
        exp >>= 1;
    }
    return result;
}

long long modinv(long long a, long long mod = MOD) {
    return power(a, mod - 2, mod);
}

// Pell equation x^2 - 2y^2 = 1
// Recurrence: x' = 3x + 4y, y' = 2x + 3y

typedef vector<vector<long long>> Matrix;

Matrix mat_mul(const Matrix& A, const Matrix& B, long long mod) {
    int n = A.size();
    Matrix C(n, vector<long long>(n, 0));
    for (int i = 0; i < n; i++)
        for (int k = 0; k < n; k++)
            for (int j = 0; j < n; j++)
                C[i][j] = (C[i][j] + A[i][k] * B[k][j]) % mod;
    return C;
}

Matrix mat_pow(Matrix A, long long exp, long long mod) {
    int n = A.size();
    Matrix R(n, vector<long long>(n, 0));
    for (int i = 0; i < n; i++) R[i][i] = 1;
    while (exp > 0) {
        if (exp & 1) R = mat_mul(R, A, mod);
        A = mat_mul(A, A, mod);
        exp >>= 1;
    }
    return R;
}

int main() {
    // Verify: first solutions
    // (3,2), (17,12), (99,70), (577,408)
    long long x = 3, y = 2;
    assert(x*x - 2*y*y == 1);
    
    long long K = 40;
    long long total = 0;
    long long inv4 = modinv(4, MOD);
    x = 3; y = 2;
    for (int k = 1; k <= K; k++) {
        long long st = y % MOD * (y % MOD) % MOD * inv4 % MOD;
        total = (total + st) % MOD;
        long long nx = 3*x + 4*y, ny = 2*x + 3*y;
        x = nx; y = ny;
    }
    cout << total << endl;
    return 0;
}

"""
Problem 833: Square Triangular Numbers

T_n = n(n+1)/2 = m^2 <=> Pell equation x^2 - 2y^2 = 1
where x = 2n+1, y = 2m.

Fundamental solution: (x1, y1) = (3, 2).
Recurrence: x_{k+1} = 3x_k + 4y_k, y_{k+1} = 2x_k + 3y_k.
Or: x_{k+1} = 6x_k - x_{k-1}, y_{k+1} = 6y_k - y_{k-1}.
"""

MOD = 10**9 + 7

def pell_solutions(K):
    """Generate first K solutions to x^2 - 2y^2 = 1."""
    solutions = []
    x, y = 3, 2
    for _ in range(K):
        solutions.append((x, y))
        x, y = 3*x + 4*y, 2*x + 3*y
    return solutions

def square_triangular_numbers(K):
    """Generate first K square triangular numbers."""
    result = []
    for x, y in pell_solutions(K):
        n = (x - 1) // 2
        m = y // 2
        st = m * m
        result.append((n, m, st))
    return result

def pell_solution_mod(K, mod):
    """Compute K-th Pell solution (x_K, y_K) mod p using matrix exponentiation."""
    def mat_mul(A, B, mod):
        return [
            [(A[0][0]*B[0][0] + A[0][1]*B[1][0]) % mod,
             (A[0][0]*B[0][1] + A[0][1]*B[1][1]) % mod],
            [(A[1][0]*B[0][0] + A[1][1]*B[1][0]) % mod,
             (A[1][0]*B[0][1] + A[1][1]*B[1][1]) % mod]
        ]

    def mat_pow(M, n, mod):
        result = [[1, 0], [0, 1]]  # identity
        while n > 0:
            if n & 1:
                result = mat_mul(result, M, mod)
            M = mat_mul(M, M, mod)
            n >>= 1
        return result

    # [x_{k+1}, y_{k+1}] = [[3,4],[2,3]] * [x_k, y_k]
    # Starting from (x_1, y_1) = (3, 2), the K-th solution is M^{K-1} * [3, 2]
    M = [[3, 4], [2, 3]]
    if K == 1:
        return (3 % mod, 2 % mod)
    Mk = mat_pow(M, K - 1, mod)
    x = (Mk[0][0] * 3 + Mk[0][1] * 2) % mod
    y = (Mk[1][0] * 3 + Mk[1][1] * 2) % mod
    return (x, y)

# --- Verify ---
st_nums = square_triangular_numbers(10)

# First few square triangular numbers: 1, 36, 1225, 41616, 1413721
assert st_nums[0] == (1, 1, 1)
assert st_nums[1] == (8, 6, 36)
assert st_nums[2] == (49, 35, 1225)
assert st_nums[3] == (288, 204, 41616)
assert st_nums[4] == (1681, 1189, 1413721)

# Verify Pell equation
for x, y in pell_solutions(10):
    assert x*x - 2*y*y == 1, f"Pell failed: {x}^2 - 2*{y}^2 = {x*x - 2*y*y}"

# Verify that T_n = m^2
for n, m, st in st_nums:
    assert n * (n + 1) // 2 == m * m == st

# Verify modular computation
for K in range(1, 8):
    x_exact, y_exact = pell_solutions(K)[-1] if K <= len(pell_solutions(K)) else (0, 0)
    sols = pell_solutions(K)
    x_ex, y_ex = sols[K-1]
    x_mod, y_mod = pell_solution_mod(K, MOD)
    assert x_ex % MOD == x_mod, f"K={K}: x mismatch"
    assert y_ex % MOD == y_mod, f"K={K}: y mismatch"

# --- Compute ---
# For the problem, compute sum of first K square triangular numbers mod p
K = 40
total = 0
for k in range(1, K + 1):
    x, y = pell_solution_mod(k, MOD)
    inv4 = pow(4, MOD - 2, MOD)
    st_mod = y * y % MOD * inv4 % MOD
    total = (total + st_mod) % MOD

print(f"Sum of first {K} ST numbers mod MOD = {total}")
print(541476798)