messy_dp_day18_upping.py

from math import sqrt
import time

SQRT_2 = sqrt(2)
INFINITY = float('inf')

def get_val(row, i):
	return row[i] if i >=0 and i < len(row) else '.'

def next_row(row):
	N = len(row)
	return [row[1]] + ['^' if row[i-1] != row[i+1] else '.' for i in xrange(1, N-1)] + [row[N-2]]

# distance represented by (x,y) = x + sqrt(2) * y
# None returns "infinity"
def get_dist(tup):
	if tup is None:
		return INFINITY
	integer, sqrt2factor = tup
	return integer + SQRT_2 * sqrt2factor

# add (x1,y1), (x2,y2)
def add_tups(tup1, tup2):
	return (tup1[0] + tup2[0], tup1[1] + tup2[1])

def solve(row, n_rows):
	N = len(row)
	# keep track of minimum distance to reach each tile in the 2 previous rows
	best_last_2 = [[(0, 0) if c == '.' and j == 0 else None for c in row] for j in xrange(2)]
	for it in xrange(1, n_rows):
		if it % 10000 == 0:
			print it
		row = next_row(row)
		blr = best_last_2[1 - it%2] # best distance to each tile in last row
		blr2 = best_last_2[it%2] # best distance to each tile 2 rows ago
		best_cur = [add_tups(blr2[i], (2, 0)) if blr2[i] is not None and row[i] == '.' else None for i in xrange(N)]
		for i in xrange(N):
			if row[i] == '^':
				continue
			if i > 0 and blr[i-1] is not None:
				move_from_left = add_tups(blr[i-1], (0, 1))
				if get_dist(move_from_left) < get_dist(best_cur[i]):
					best_cur[i] = move_from_left
			if blr[i] is not None:
				move_from_top = add_tups(blr[i], (1, 0))
				if get_dist(move_from_top) < get_dist(best_cur[i]):
					best_cur[i] = move_from_top
			if i < N-1 and blr[i+1] is not None:
				move_from_right = add_tups(blr[i+1], (0, 1))
				if get_dist(move_from_right) < get_dist(best_cur[i]):
					best_cur[i] = move_from_right
		best_last_2[it%2] = best_cur # replace best moves for 2 rows ago with current row
	return sorted([(get_dist(b), b) for b in best_cur])

start = time.time()
print solve('.^^^^^.^^^..^^^^^...^.^..^^^.^^....^.^...^^^...^^^^..^...^...^^.^.^.......^..^^...^.^.^^..^^^^^...^.', 400000)
print time.time() - start