MMA CTF 2015 Lock Pattern - 20

mmactf_ppc40_1.py

from functools import lru_cache
from math import sqrt
from itertools import product
from collections import defaultdict, deque

size = 3

def gcd(x, y):
	if (y == 0): return x
	else: return gcd(y, x%y)
def sqr(x): return x*x
def popcnt(x): return bin(x).count('1')
def pack(x, y): return (1 << x*size+y)

# get points between two point
@lru_cache(maxsize=None)
def getBetween(p1, p2):
	if (p1 == p2): return []
	res = []
	x1, y1 = p1
	x2, y2 = p2
	# vector between two points
	dx, dy = x2 - x1, y2 - y1
	g = gcd(abs(dx), abs(dy))
	dx, dy = dx // g, dy // g
	# check for points lie between them
	x, y = x1 + dx, y1 + dy
	while not ((x == x2) and (y == y2)):
		res.append((x, y))
		x += dx
		y += dy
	return res

count = 0
dp = defaultdict(int)
q = deque()
# init base cases
for x, y in product(range(size), repeat=2):
	dp[(pack(x, y), x, y)] = 1
	q.append((pack(x, y), x, y))
enqueued = set()
while (q):
	mask, x, y = q.popleft()
	# check second condition
	if (popcnt(mask) >= 4): count += dp[(mask, x, y)]
	for i, j in product(range(size), repeat=2):
		if (mask & pack(i, j)): continue
		# check third condition
		req = 0
		for u, v in getBetween((x, y), (i, j)):
			req = req | pack(u, v)
		if ((req & mask) != req): continue
		dp[(mask | pack(i, j), i, j)] += dp[mask, x, y]
		if ((mask | pack(i, j), i, j) not in enqueued):
			q.append((mask | pack(i, j), i, j))
			enqueued.add((mask | pack(i, j), i, j))
print(count)