Add Euler project problem 124 solution

project_euler/problem_124/sol1.py

@@ -0,0 +1,109 @@
+"""
+Project Euler Problem 124: https://projecteuler.net/problem=124
+
+Ordered Radicals
+
+"""
+
+from numpy import sqrt
+
+
+def generate_primes(n: int) -> list[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 generate_n(factors: list[int], n_max: int, n: int, res: set[int]):
+    """
+    Generates all numbers n that can be constructed out of 'factors', with any
+    multiplicity, but that do no exceed 'n_max'.
+
+    >>> generate_n([2], 10, 1, set())
+    """
+
+    if len(factors) == 0:
+        return
+    fac = factors[0]
+    factors_new = factors[1:]
+    while n <= n_max:
+        generate_n(factors_new, n_max, n, res)
+        res.add(n)
+        n *= fac
+    return
+
+
+def generate_rads(
+    factors_all: list[int], n_max: int, n: int, res: dict, factors_prev: list[int]
+):
+    """
+    Generates all rads and associated factors, e.g., rad = factor_1 * ... * factor_k.
+    Output is stored in 'res' dict argument.
+
+    >>> generate_rads([2], 10, 1, {}, [])
+    """
+
+    for i in range(len(factors_all)):
+        f = factors_all[i]
+        n_new = n * f
+        if n_new > n_max:
+            return
+        # factors_new = factors_prev + [f]
+        factors_new = [*factors_prev, f]
+        res[n_new] = factors_new
+        generate_rads(factors_all[(i + 1) :], n_max, n_new, res, factors_new)
+    return
+
+
+def solution(n_max: int = 100000, k: int = 10000) -> int:
+    """
+    Loops over sorted 'rads' and generates all numbers 'n' for rad.
+    Keeps track of total number of n, and when k falls inside some rad,
+    it sorts all 'n' for it and picks up associated n.
+
+    >>> solution(10, 6)
+    9
+    >>> solution(10, 9)
+    7
+    """
+
+    if k == 1:
+        return 1
+
+    primes = generate_primes(n_max)
+    tot = 1
+    rads_d: dict[int, list[int]] = {}
+    factor_prev: list[int] = []
+    generate_rads(primes, n_max, 1, rads_d, factor_prev)
+    rads = sorted(rads_d)
+
+    for r in rads:
+        facts = rads_d[r]
+        res: set[int] = set()
+        generate_n(facts, n_max, r, res)
+        res_len = len(res)
+        if tot + res_len >= k:
+            return sorted(res)[k - tot - 1]
+        tot += res_len
+    return -1
+
+
+if __name__ == "__main__":
+    print(f"{solution() = }")