diff --git a/python/e313.py b/python/e313.py new file mode 100644 index 0000000..f3e592c --- /dev/null +++ b/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