Problem 1008: Functional Inverse

Problem 1008 source code

C++

#include <algorithm>
#include <array>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <stdexcept>
#include <string>
#include <thread>
#include <vector>

namespace {

using u64 = std::uint64_t;
using Polynomial = std::array<u64, 10>;

constexpr u64 MOD = 1'000'000'007;
constexpr int TARGET = 10'000'000;
constexpr int BATCH_SIZE = 32'768;

u64 add(const u64 a, const u64 b) {
    const u64 sum = a + b;
    return sum >= MOD ? sum - MOD : sum;
}

u64 subtract(const u64 a, const u64 b) {
    return a >= b ? a - b : a + MOD - b;
}

u64 multiply(const u64 a, const u64 b) {
    return a * b % MOD;
}

u64 power(u64 base, u64 exponent) {
    u64 result = 1;
    while (exponent != 0) {
        if ((exponent & 1) != 0) result = multiply(result, base);
        base = multiply(base, base);
        exponent >>= 1;
    }
    return result;
}

struct Block {
    Polynomial product{1};
    Polynomial sum{};
    u64 factorial = 1;
};

template <bool FullSum>
Block calculate_block(const int begin, const int end) {
    Block result;
    std::vector<u64> inverse(BATCH_SIZE + 1);

    for (int first = begin; first < end; first += BATCH_SIZE) {
        const int count = std::min(BATCH_SIZE, end - first);
        inverse[0] = 1;
        for (int i = 1; i <= count; ++i) {
            const u64 k = static_cast<u64>(first + i - 1);
            inverse[i] = multiply(inverse[i - 1], multiply(k, 2 * k - 1));
        }
        u64 suffix_inverse = power(inverse[count], MOD - 2);
        for (int i = count; i >= 1; --i) {
            const u64 k = static_cast<u64>(first + i - 1);
            const u64 weight = multiply(suffix_inverse, inverse[i - 1]);
            suffix_inverse = multiply(suffix_inverse, multiply(k, 2 * k - 1));
            inverse[i] = weight;
        }

        for (int i = 1; i <= count; ++i) {
            const u64 k = static_cast<u64>(first + i - 1);
            const u64 weight = inverse[i];
            const u64 inverse_k = multiply(2 * k - 1, weight);
            const u64 inverse_square = multiply(inverse_k, inverse_k);
            result.factorial = multiply(result.factorial, k);

            // P_N/x = sum A_(k-1)/(k(2k-1)), A_k = product_(j<=k)(1-x/j^2).
            if constexpr (FullSum) {
                for (int degree = 9; degree >= 1; --degree) {
                    result.sum[degree] = add(result.sum[degree],
                                             multiply(weight, result.product[degree]));
                    result.product[degree] = subtract(result.product[degree],
                        multiply(inverse_square, result.product[degree - 1]));
                }
                result.sum[0] = add(result.sum[0], weight);
            } else {
                result.sum[9] = add(result.sum[9], multiply(weight, result.product[9]));
                for (int degree = 9; degree >= 1; --degree) {
                    result.product[degree] = subtract(result.product[degree],
                        multiply(inverse_square, result.product[degree - 1]));
                }
            }
        }
    }
    return result;
}

Polynomial product(const Polynomial& a, const Polynomial& b) {
    Polynomial result{};
    for (int i = 0; i <= 9; ++i) {
        for (int j = 0; j <= i; ++j) {
            result[i] = add(result[i], multiply(a[j], b[i - j]));
        }
    }
    return result;
}

u64 solve(const int n, unsigned thread_count) {
    thread_count = std::max(1U, std::min(thread_count, static_cast<unsigned>(n)));
    Block total;
    if (thread_count == 1) {
        total = calculate_block<false>(1, n + 1);
    } else {
        std::vector<Block> blocks(thread_count);
        std::vector<std::thread> workers;
        for (unsigned index = 0; index < thread_count; ++index) {
            const int begin = 1 + static_cast<int>(static_cast<u64>(n) * index / thread_count);
            const int end = 1 + static_cast<int>(static_cast<u64>(n) * (index + 1) / thread_count);
            workers.emplace_back([&, index, begin, end] {
                blocks[index] = calculate_block<true>(begin, end);
            });
        }
        for (auto& worker : workers) worker.join();
        for (const Block& block : blocks) {
            const Polynomial contribution = product(total.product, block.sum);
            total.sum[9] = add(total.sum[9], contribution[9]);
            total.product = product(total.product, block.product);
            total.factorial = multiply(total.factorial, block.factorial);
        }
    }

    const u64 factorial_square = multiply(total.factorial, total.factorial);
    u64 leading = multiply(static_cast<u64>(n),
        power(multiply(2ULL * n - 1, factorial_square), MOD - 2));
    if (n % 2 == 0) leading = subtract(0, leading);
    if (leading == 1) return total.sum[9];

    // Otherwise the minimum monic polynomial is P_N + x*product_(k<=N)(x-k^2).
    const u64 scale = n % 2 == 0 ? factorial_square : subtract(0, factorial_square);
    return add(total.sum[9], multiply(scale, total.product[9]));
}

std::vector<u64> append_root(const std::vector<u64>& polynomial, const u64 root) {
    std::vector<u64> result(polynomial.size() + 1);
    for (std::size_t i = 0; i < polynomial.size(); ++i) {
        result[i] = subtract(result[i], multiply(root, polynomial[i]));
        result[i + 1] = add(result[i + 1], polynomial[i]);
    }
    return result;
}

u64 interpolate_directly(const int n) {
    std::vector<u64> polynomial(static_cast<std::size_t>(n + 2));
    for (int i = 1; i <= n; ++i) {
        std::vector<u64> basis{1};
        u64 denominator = 1;
        for (int j = 0; j <= n; ++j) {
            if (j == i) continue;
            basis = append_root(basis, static_cast<u64>(j) * j);
            denominator = multiply(denominator,
                subtract(static_cast<u64>(i) * i, static_cast<u64>(j) * j));
        }
        const u64 scale = multiply(static_cast<u64>(i), power(denominator, MOD - 2));
        for (std::size_t degree = 0; degree < basis.size(); ++degree) {
            polynomial[degree] = add(polynomial[degree], multiply(scale, basis[degree]));
        }
    }
    if (polynomial[n] != 1) {
        std::vector<u64> vanishing{1};
        for (int i = 0; i <= n; ++i) {
            vanishing = append_root(vanishing, static_cast<u64>(i) * i);
        }
        for (std::size_t degree = 0; degree < vanishing.size(); ++degree) {
            polynomial[degree] = add(polynomial[degree], vanishing[degree]);
        }
    }
    return polynomial.size() > 10 ? polynomial[10] : 0;
}

void run_tests() {
    for (int n = 1; n <= 40; ++n) {
        const u64 expected = interpolate_directly(n);
        if (solve(n, 1) != expected || solve(n, 3) != expected) {
            throw std::runtime_error("Interpolation check failed for n=" + std::to_string(n));
        }
    }
    for (const int n : {BATCH_SIZE - 1, BATCH_SIZE, BATCH_SIZE + 1, 2 * BATCH_SIZE + 3}) {
        if (solve(n, 1) != solve(n, 3)) {
            throw std::runtime_error("Batch boundary check failed for n=" + std::to_string(n));
        }
    }
    std::cout << "All checks passed.\n";
}

}  // namespace

int main(int argc, char* argv[]) {
    try {
        unsigned thread_count = std::min(16U, std::max(1U, std::thread::hardware_concurrency()));
        if (argc == 2 && std::string(argv[1]) == "--self-test") {
            run_tests();
            return EXIT_SUCCESS;
        }
        if (argc == 3 && std::string(argv[1]) == "--threads") {
            const std::string value = argv[2];
            if (value.empty() || value.find_first_not_of("0123456789") != std::string::npos) {
                throw std::invalid_argument("Thread count must be an integer from 1 to 128.");
            }
            const unsigned long parsed = std::stoul(value);
            if (parsed < 1 || parsed > 128) {
                throw std::invalid_argument("Thread count must be an integer from 1 to 128.");
            }
            thread_count = static_cast<unsigned>(parsed);
        } else if (argc != 1) {
            throw std::invalid_argument("Usage: Euler1008 [--threads COUNT | --self-test]");
        }
        std::cout << solve(TARGET, thread_count) << '\n';
    } catch (const std::exception& error) {
        std::cerr << error.what() << '\n';
        return EXIT_FAILURE;
    }
    return EXIT_SUCCESS;
}

Python

#!/usr/bin/env python3
"""Project Euler Problem 1008 - Functional Inverse."""

from __future__ import annotations

MOD = 1_000_000_007
TARGET = 10_000_000
WANTED_DEGREE = 9  # P_N has a factor x, so x^10 comes from degree 9.
BATCH_SIZE = 32_768


def target_coefficient(n: int, batch_size: int = BATCH_SIZE) -> int:
    """Return the x^10 coefficient in the (n, MOD)-functional inverse.

    If A_k(x) = product(1 - x/j^2, j=1..k), Newton interpolation gives

        P_N(x) / x = sum(A_(k-1)(x) / (k(2k-1)), k=1..N).

    Only coefficients through degree nine are retained.  A batch inversion
    supplies every 1/(k(2k-1)) with one modular exponentiation per batch.
    """

    if n < 1 or 2 * n - 1 >= MOD:
        raise ValueError("this implementation requires 1 <= n and 2n-1 < MOD")
    if batch_size < 1:
        raise ValueError("batch_size must be positive")

    # coefficients[d] is [x^d] A_k for the current k.
    coefficients = [0] * (WANTED_DEGREE + 1)
    coefficients[0] = 1
    interpolation_coefficient = 0
    factorial = 1
    prefix = [1] * (min(batch_size, n) + 1)

    for first in range(1, n + 1, batch_size):
        count = min(batch_size, n - first + 1)

        # Prefix products of v_k = k(2k-1).
        prefix[0] = 1
        for index in range(1, count + 1):
            k = first + index - 1
            value = k * (2 * k - 1) % MOD
            prefix[index] = prefix[index - 1] * value % MOD

        # Recover every v_k^-1 by walking backward from the inverse of the
        # complete product.  Reuse prefix[] to store the resulting inverses.
        suffix_inverse = pow(prefix[count], MOD - 2, MOD)
        for index in range(count, 0, -1):
            k = first + index - 1
            value = k * (2 * k - 1) % MOD
            weight = suffix_inverse * prefix[index - 1] % MOD
            suffix_inverse = suffix_inverse * value % MOD
            prefix[index] = weight

        for index in range(1, count + 1):
            k = first + index - 1
            weight = prefix[index]  # 1/(k(2k-1))
            inverse_k = (2 * k - 1) * weight % MOD
            inverse_square = inverse_k * inverse_k % MOD

            # The summand uses A_(k-1), so accumulate before updating A.
            interpolation_coefficient = (
                interpolation_coefficient
                + weight * coefficients[WANTED_DEGREE]
            ) % MOD
            for degree in range(WANTED_DEGREE, 0, -1):
                coefficients[degree] = (
                    coefficients[degree]
                    - inverse_square * coefficients[degree - 1]
                ) % MOD

            factorial = factorial * k % MOD

    factorial_square = factorial * factorial % MOD
    leading = n * pow(
        (2 * n - 1) * factorial_square % MOD, MOD - 2, MOD
    ) % MOD
    if n % 2 == 0:
        leading = -leading % MOD

    # P_N is already the minimum monic polynomial only when its leading
    # coefficient is one.  Otherwise add the monic vanishing polynomial
    # x*product(x-k^2, k=1..N).
    if leading == 1:
        return interpolation_coefficient

    vanishing_scale = factorial_square if n % 2 == 0 else -factorial_square % MOD
    return (
        interpolation_coefficient
        + vanishing_scale * coefficients[WANTED_DEGREE]
    ) % MOD


def append_root(polynomial: list[int], root: int) -> list[int]:
    """Return polynomial * (x-root), with coefficients modulo MOD."""

    result = [0] * (len(polynomial) + 1)
    for degree, coefficient in enumerate(polynomial):
        result[degree] = (result[degree] - root * coefficient) % MOD
        result[degree + 1] = (result[degree + 1] + coefficient) % MOD
    return result


def interpolate_directly(n: int) -> int:
    """Build the whole small polynomial by Lagrange interpolation.

    This cubic-time reference is intentionally independent of the truncated
    Newton recurrence and is used only by the checkpoints below.
    """

    polynomial = [0] * (n + 2)
    for i in range(1, n + 1):
        basis = [1]
        denominator = 1
        node_i = i * i % MOD
        for j in range(n + 1):
            if j == i:
                continue
            node_j = j * j % MOD
            basis = append_root(basis, node_j)
            denominator = denominator * (node_i - node_j) % MOD

        scale = i * pow(denominator, MOD - 2, MOD) % MOD
        for degree, coefficient in enumerate(basis):
            polynomial[degree] = (
                polynomial[degree] + scale * coefficient
            ) % MOD

    if polynomial[n] != 1:
        vanishing = [1]
        for i in range(n + 1):
            vanishing = append_root(vanishing, i * i % MOD)
        for degree, coefficient in enumerate(vanishing):
            polynomial[degree] = (polynomial[degree] + coefficient) % MOD

    return polynomial[10] if len(polynomial) > 10 else 0


def run_checkpoints() -> None:
    for n in range(1, 41):
        expected = interpolate_directly(n)
        assert target_coefficient(n) == expected, f"interpolation check for n={n}"
        assert target_coefficient(n, 7) == expected, f"batch check for n={n}"


def main() -> None:
    run_checkpoints()
    print(target_coefficient(TARGET))


if __name__ == "__main__":
    main()

Java

public class Euler1008 {
    private static final long MOD = 1_000_000_007L;
    private static final int TARGET = 10_000_000;
    // P_N has a factor x, so its x^10 coefficient comes from degree nine.
    private static final int WANTED_DEGREE = 9;
    private static final int BATCH_SIZE = 32_768;

    private static long add(long lhs, long rhs) {
        long sum = lhs + rhs;
        return sum >= MOD ? sum - MOD : sum;
    }

    private static long subtract(long lhs, long rhs) {
        return lhs >= rhs ? lhs - rhs : lhs + MOD - rhs;
    }

    private static long multiply(long lhs, long rhs) {
        // Both operands are below MOD, so their product fits in signed long.
        return lhs * rhs % MOD;
    }

    private static long modPow(long base, long exponent) {
        long result = 1;
        while (exponent != 0) {
            if ((exponent & 1L) != 0) {
                result = multiply(result, base);
            }
            base = multiply(base, base);
            exponent >>= 1;
        }
        return result;
    }

    private static long targetCoefficient(int n) {
        return targetCoefficient(n, BATCH_SIZE);
    }

    private static long targetCoefficient(int n, int batchSize) {
        if (n < 1 || 2L * n - 1 >= MOD) {
            throw new IllegalArgumentException(
                    "this implementation requires 1 <= n and 2n-1 < MOD");
        }
        if (batchSize < 1) {
            throw new IllegalArgumentException("batchSize must be positive");
        }

        // coefficients[d] is [x^d] A_k, where
        // A_k(x) = product(1 - x/j^2, j=1..k).
        long[] coefficients = new long[WANTED_DEGREE + 1];
        coefficients[0] = 1;
        long interpolationCoefficient = 0;
        long factorial = 1;
        long[] prefix = new long[Math.min(batchSize, n) + 1];

        for (int first = 1; first <= n; first += batchSize) {
            int count = Math.min(batchSize, n - first + 1);

            // Prefix products of v_k = k(2k-1).
            prefix[0] = 1;
            for (int index = 1; index <= count; ++index) {
                long k = (long) first + index - 1;
                long value = multiply(k % MOD, (2 * k - 1) % MOD);
                prefix[index] = multiply(prefix[index - 1], value);
            }

            // One Fermat inverse followed by a backward pass supplies every
            // 1/v_k in the batch.  Store those weights back into prefix[].
            long suffixInverse = modPow(prefix[count], MOD - 2);
            for (int index = count; index >= 1; --index) {
                long k = (long) first + index - 1;
                long value = multiply(k % MOD, (2 * k - 1) % MOD);
                long weight = multiply(suffixInverse, prefix[index - 1]);
                suffixInverse = multiply(suffixInverse, value);
                prefix[index] = weight;
            }

            for (int index = 1; index <= count; ++index) {
                long k = (long) first + index - 1;
                long weight = prefix[index]; // 1/(k(2k-1))
                long inverseK = multiply((2 * k - 1) % MOD, weight);
                long inverseSquare = multiply(inverseK, inverseK);

                // Newton interpolation gives
                // P_N/x = sum A_(k-1)/(k(2k-1)).  Accumulate the summand
                // before changing A_(k-1) into A_k.
                interpolationCoefficient = add(
                        interpolationCoefficient,
                        multiply(weight, coefficients[WANTED_DEGREE]));
                for (int degree = WANTED_DEGREE; degree >= 1; --degree) {
                    coefficients[degree] = subtract(
                            coefficients[degree],
                            multiply(inverseSquare, coefficients[degree - 1]));
                }

                factorial = multiply(factorial, k % MOD);
            }
        }

        long factorialSquare = multiply(factorial, factorial);
        long leading = multiply(
                n,
                modPow(multiply((2L * n - 1) % MOD, factorialSquare), MOD - 2));
        if ((n & 1) == 0) {
            leading = subtract(0, leading);
        }

        // If P_N is not monic, the lowest-degree monic solution is obtained
        // by adding x*product(x-k^2, k=1..N).
        if (leading == 1) {
            return interpolationCoefficient;
        }

        long vanishingScale = (n & 1) == 0
                ? factorialSquare
                : subtract(0, factorialSquare);
        return add(
                interpolationCoefficient,
                multiply(vanishingScale, coefficients[WANTED_DEGREE]));
    }

    private static long[] appendRoot(long[] polynomial, long root) {
        long[] result = new long[polynomial.length + 1];
        for (int degree = 0; degree < polynomial.length; ++degree) {
            result[degree] = subtract(
                    result[degree], multiply(root, polynomial[degree]));
            result[degree + 1] = add(result[degree + 1], polynomial[degree]);
        }
        return result;
    }

    // Independent cubic-time Lagrange interpolation for small checkpoints.
    private static long interpolateDirectly(int n) {
        long[] polynomial = new long[n + 2];

        for (int i = 1; i <= n; ++i) {
            long[] basis = {1};
            long denominator = 1;
            long nodeI = (long) i * i % MOD;

            for (int j = 0; j <= n; ++j) {
                if (j == i) {
                    continue;
                }
                long nodeJ = (long) j * j % MOD;
                basis = appendRoot(basis, nodeJ);
                denominator = multiply(denominator, subtract(nodeI, nodeJ));
            }

            long scale = multiply(i, modPow(denominator, MOD - 2));
            for (int degree = 0; degree < basis.length; ++degree) {
                polynomial[degree] = add(
                        polynomial[degree], multiply(scale, basis[degree]));
            }
        }

        if (polynomial[n] != 1) {
            long[] vanishing = {1};
            for (int i = 0; i <= n; ++i) {
                vanishing = appendRoot(vanishing, (long) i * i % MOD);
            }
            for (int degree = 0; degree < vanishing.length; ++degree) {
                polynomial[degree] = add(polynomial[degree], vanishing[degree]);
            }
        }

        return polynomial.length > 10 ? polynomial[10] : 0;
    }

    private static void requireCheckpoint(boolean condition, String description) {
        if (!condition) {
            throw new AssertionError("checkpoint failed: " + description);
        }
    }

    private static void runCheckpoints() {
        for (int n = 1; n <= 40; ++n) {
            long expected = interpolateDirectly(n);
            requireCheckpoint(
                    targetCoefficient(n) == expected,
                    "interpolation check for n=" + n);
            requireCheckpoint(
                    targetCoefficient(n, 7) == expected,
                    "batch check for n=" + n);
        }
    }

    public static void main(String[] args) {
        runCheckpoints();
        System.out.println(targetCoefficient(TARGET));
    }
}