sol1.py

"""
Project Euler Problem 95: https://projecteuler.net/problem=95

Amicable Chains

Solution is doing the following:
- Get relevant prime numbers
- Iterate over product combination of prime numbers to generate all non-prime
numbers up to max number, by keeping track of prime factors
- Calculate the sum of factors for each number
- Iterate over found some factors to find longest chain

>>> solution(200000)
12496

"""

from numpy import sqrt


def sum_primes(factor_d, num):
    """
    Calculates the sum of factors from all prime exponents.

    >>> sum_primes({2: 1, 3: 1}, 6)
    6
    """
    tot = 1
    for p in factor_d:
        comp = 0
        ex_factor = 1
        for _ in range(factor_d[p] + 1):
            comp += ex_factor
            ex_factor *= p
        tot *= comp
    return tot - num


def generate_primes(n: int):
    """
    Calculates the list of primes up to and including n.

    >>> generate_primes(6)
    [2, 3, 5]
    """
    primes = [True] * (n + 1)
    primes[0] = primes[1] = False
    for i in range(2, int(sqrt(n + 1)) + 1):
        if primes[i]:
            j = i * i
            while j <= n:
                primes[j] = False
                j += i
    primes_list = []
    for i in range(2, len(primes)):
        if primes[i]:
            primes_list += [i]
    return primes_list


def multiply(chain, primes, prime, prev_n, n_max, prev_sum, primes_d):
    """
    Run over all prime combinations to generate non-prime numbers.

    >>> multiply([None] * 3, {2}, 2, 1, 2, 0, {})
    """

    number = prev_n * prime
    primes_d[prime] = primes_d.get(prime, 0) + 1
    if prev_n % prime != 0:
        new_sum = prev_sum * (prime + 1) + prev_n
    else:
        new_sum = sum_primes(primes_d, number)
    chain[number] = new_sum
    for p in primes:
        if p >= prime:
            number_n = p * number
            if number_n > n_max:
                break
            multiply(chain, primes, p, number, n_max, new_sum, primes_d.copy())


def find_longest_chain(chain, n_max):
    """
    Finds the smallest element and length of longest chain

    >>> find_longest_chain([0, 0, 0, 0, 0, 0, 6], 6)
    (6, 1)
    """

    length_max = 0
    elem_max = 0
    for i in range(2, len(chain)):
        start = i
        length = 1
        el = chain[i]
        visited = {i}
        while el > 1 and el <= n_max and el not in visited:
            length += 1
            visited.add(el)
            el = chain[el]

        if el == start and length > length_max:
            length_max = length
            elem_max = start

    return elem_max, length_max


def solution(n_max: int = 1000000) -> int:
    """
    Runs the calculation for numbers <= n_max.

    >>> solution(10)
    6
    """

    primes = generate_primes(n_max)
    chain = [0] * (n_max + 1)
    for p in primes:
        if p * p > n_max:
            break
        multiply(chain, primes, p, 1, n_max, 0, {})

    chain_start, _ = find_longest_chain(chain, n_max)
    return chain_start


if __name__ == "__main__":
    print(f"{solution() = }")