Solve problem 313

python/e313.py

@@ -0,0 +1,99 @@
+from lib_prime import primes, is_prime
+from dataclasses import dataclass
+from heapq import heappop, heappush
+from collections import namedtuple
+from collections import defaultdict
+from math import sqrt
+
+P = namedtuple("P", ["x", "y"])
+
+State = namedtuple("State", ["f", "g", "red", "space"])
+
+
+@dataclass(order=True)
+class State:
+    g: int
+    red: P
+    space: P
+
+
+def s(m: int, n: int) -> int:
+    """What am I doing? It was necessary to find the pattern."""
+
+    def h(s: State) -> int:
+        return abs(m - s.space.x) + abs(n - s.space.y)
+
+    def f(s: State) -> int:
+        return s.g + h(s)
+
+    def done(s: State) -> bool:
+        return s.red.x == m - 1 and s.red.y == n - 1
+
+    # assert m < n, then we could just take result times 2
+    init = State(0, P(0, 0), P(m - 1, n - 1))
+    states = [(0, init)]
+    g_costs = {(init.red, init.space): init.g}
+
+    while states:
+        _, s = heappop(states)
+        if done(s):
+            return s.g
+
+        for dx, dy in [(-1, 0), (0, 1), (1, 0), (0, -1)]:
+            nx, ny = s.space.x + dx, s.space.y + dy
+            if not (0 <= nx < m and 0 <= ny < n):
+                continue
+            nstate = State(s.g + 1, s.red, P(nx, ny))
+            if s.red.x == nx and s.red.y == ny:
+                nstate.red = P(s.space.x, s.space.y)
+            if (
+                not (nstate.red, nstate.space) in g_costs
+                or g_costs[(nstate.red, nstate.space)] > nstate.g
+            ):
+                g_costs[nstate.red, nstate.space] = nstate.g
+                heappush(states, (f(nstate), nstate))
+    return -1
+
+
+def solve_naiv(p_max: int) -> int:
+    t = 0
+    p2_max = p_max * p_max
+    ps = primes(p_max)
+    p2s = set([p * p for p in ps])
+    for m in range(2, p2_max):
+        for n in range(m, p2_max):
+            r = s(m, n)
+            if r in p2s:
+                t += 2
+            if r > p2_max:
+                break
+    return t
+
+
+def solve(p_max: int) -> int:
+    t = 2
+    r_base = 0
+    delta = 1
+    three = 3
+    for p in range(3, p_max + 1, 2):
+        if is_prime(p):
+            t += r_base * 2
+        r_base += delta
+        three -= 1
+        if three == 0:
+            delta += 1
+            three = 3
+    return t
+
+
+def euler_313():
+    assert solve_naiv(5) == 4
+    assert solve(20) == 86
+    assert solve(100) == 5482
+    return solve(10**6)
+
+
+if __name__ == "__main__":
+    solution = euler_313()
+    print(f"e313.py: {solution}")
+    # assert solution == 2057774861813004