Solve problem 315

python/e315.py

@@ -0,0 +1,111 @@
+from lib_prime import is_prime_rabin_miller
+from typing import Optional
+from functools import cache
+
+
+def compute_deltas() -> list[list[int]]:
+    S = [
+        set([0, 1, 2, 4, 5, 6]),  # 0
+        set([2, 5]),  # 1
+        set([0, 2, 3, 4, 6]),  # 2
+        set([0, 2, 3, 5, 6]),  # 3
+        set([1, 2, 3, 5]),  # 4
+        set([0, 1, 3, 5, 6]),  # 5
+        set([0, 1, 3, 4, 5, 6]),  # 6
+        set([0, 1, 2, 5]),  # 7
+        set([0, 1, 2, 3, 4, 5, 6]),  # 8
+        set([0, 1, 2, 3, 5, 6]),  # 9
+    ]
+
+    xs = [[0] * 10 for _ in range(10)]
+    for a in range(10):
+        for b in range(10):
+            xs[a][b] = len(S[a] - S[b]) + len(S[b] - S[a])
+    return xs
+
+
+NUM_TO_SEGMENTS = [6, 2, 5, 5, 4, 5, 6, 4, 7, 6]
+DELTA_SEGMENTS = compute_deltas()
+
+
+@cache
+def digital_root(n: int) -> int:
+    root = 0
+    while n != 0:
+        d = n % 10
+        n = n // 10
+        root += d
+    return root
+
+
+@cache
+def total_segments(n: int) -> int:
+    count = 0
+    n_next = 0
+    while n != 0:
+        d = n % 10
+        n = n // 10
+        count += NUM_TO_SEGMENTS[d]
+        n_next += d
+    return count
+
+
+@cache
+def transition_segments(a: int, b: int) -> int:
+    count = 0
+    while a != 0 or b != 0:
+        assert a >= b
+        da, db = a % 10, b % 10
+        if b == 0:
+            # We are just turning off the remaining digits
+            count += NUM_TO_SEGMENTS[da]
+        else:
+            count += DELTA_SEGMENTS[da][db]
+        a, b = a // 10, b // 10
+    return count
+
+
+@cache
+def sam_count(n: int) -> int:
+    count = total_segments(n)
+    count *= 2  # turn off n
+    n_next = digital_root(n)
+    if n == n_next:
+        return count
+    return count + sam_count(n_next)
+
+
+@cache
+def max_count(n_current: Optional[int], n: int) -> int:
+    if n_current is None:
+        # First number, only turn number on
+        count = total_segments(n)
+        n_next = digital_root(n)
+        assert n != n_next
+        return count + max_count(n, n_next)
+    elif n_current < 10:
+        # Last number, only turn number off
+        assert n_current == n
+        return total_segments(n)
+
+    # This is the case where we have to compute the actual delta
+    count = transition_segments(n_current, n)
+    n_next = digital_root(n)
+    return count + max_count(n, n_next)
+
+
+def euler_315():
+    assert sam_count(137) == 40
+    assert max_count(None, 137) == 30
+
+    t = 0
+    for n in range(10**7, 2 * 10**7):
+        if is_prime_rabin_miller(n):
+            t += sam_count(n) - max_count(None, n)
+    return t
+
+
+if __name__ == "__main__":
+    solution = euler_315()
+    print(f"e315.py: {solution}")
+    assert solution == 13625242