Project Euler 144

http://projecteuler.net/index.php?section=problems&id=144


反射ベクトルvは、入射ベクトルv0、鏡の法線ベクトルをnとすると、

v = v0 + k n
(v + v0)・n = 0

が成り立ちます。

from itertools import *

def unfold(f, s):
    while True:
        x, s = f(s)
        yield x

def cross_pt(pt0, v0):
    x0, y0 = pt0
    vx0, vy0 = v0
    t = -(8 * x0 * vx0 + 2 * y0 * vy0) / (4 * vx0 * vx0 + vy0 * vy0)
    return (x0 + t * vx0, y0 + t * vy0)

def reflect(pt, v0):
    x, y = pt
    v0x, v0y = v0
    nx, ny = -4 * x, -y
    k = -2 * (v0x * nx + v0y * ny) / (nx * nx + ny * ny)
    return (v0x + k * nx, v0y + k * ny)

def is_in_hole(pt):
    x, y = pt
    return y > 0 and abs(x) <= 0.01

def f(s):
    pt, v = s
    pt1 = cross_pt(pt, v)
    v1 = reflect(pt1, v)
    return (pt1, (pt1, v1))

v = (1.4, -9.6 - 10.1)
init = (cross_pt((1.4, -9.6), v), v)
print next(k for k, pt in izip(count(), unfold(f, init))
                                            if is_in_hole(pt))