GCHQ Christmas puzzle 2015

gchq-puzzle.py

# http://www.gchq.gov.uk/SiteCollectionImages/grid-shading-puzzle.jpg
# (Problem definition taken from https://github.com/matthewearl/gchq-xmas/ )
SIZE = 25
ROW_RUNS = [
    [7, 3, 1, 1, 7,],
    [1, 1, 2, 2, 1, 1,],
    [1, 3, 1, 3, 1, 1, 3, 1,],
    [1, 3, 1, 1, 6, 1, 3, 1,],
    [1, 3, 1, 5, 2, 1, 3, 1,],
    [1, 1, 2, 1, 1,],
    [7, 1, 1, 1, 1, 1, 7,],
    [3, 3,],
    [1, 2, 3, 1, 1, 3, 1, 1, 2,],
    [1, 1, 3, 2, 1, 1,],
    [4, 1, 4, 2, 1, 2,],
    [1, 1, 1, 1, 1, 4, 1, 3,],
    [2, 1, 1, 1, 2, 5,],
    [3, 2, 2, 6, 3, 1,],
    [1, 9, 1, 1, 2, 1,],
    [2, 1, 2, 2, 3, 1,],
    [3, 1, 1, 1, 1, 5, 1,],
    [1, 2, 2, 5,],
    [7, 1, 2, 1, 1, 1, 3,],
    [1, 1, 2, 1, 2, 2, 1,],
    [1, 3, 1, 4, 5, 1,],
    [1, 3, 1, 3, 10, 2,],
    [1, 3, 1, 1, 6, 6,],
    [1, 1, 2, 1, 1, 2,],
    [7, 2, 1, 2, 5,],
]
COL_RUNS = [
    [7, 2, 1, 1, 7,],
    [1, 1, 2, 2, 1, 1,],
    [1, 3, 1, 3, 1, 3, 1, 3, 1,],
    [1, 3, 1, 1, 5, 1, 3, 1,],
    [1, 3, 1, 1, 4, 1, 3, 1,],
    [1, 1, 1, 2, 1, 1,],
    [7, 1, 1, 1, 1, 1, 7,],
    [1, 1, 3,],
    [2, 1, 2, 1, 8, 2, 1,],
    [2, 2, 1, 2, 1, 1, 1, 2,],
    [1, 7, 3, 2, 1,],
    [1, 2, 3, 1, 1, 1, 1, 1,],
    [4, 1, 1, 2, 6,],
    [3, 3, 1, 1, 1, 3, 1,],
    [1, 2, 5, 2, 2,],
    [2, 2, 1, 1, 1, 1, 1, 2, 1,],
    [1, 3, 3, 2, 1, 8, 1,],
    [6, 2, 1,],
    [7, 1, 4, 1, 1, 3,],
    [1, 1, 1, 1, 4,],
    [1, 3, 1, 3, 7, 1,],
    [1, 3, 1, 1, 1, 2, 1, 1, 4,],
    [1, 3, 1, 4, 3, 3,],
    [1, 1, 2, 2, 2, 6, 1,],
    [7, 1, 3, 2, 1, 1,],
]
GIVENS = [
  (3, 3), (3, 4), (3, 12), (3, 13), (3, 21),
  (8, 6), (8, 7), (8, 10), (8, 14), (8, 15), (8, 18),
  (16, 6), (16, 11), (16, 16), (16, 20),
  (21, 3), (21, 4), (21, 9), (21, 10), (21, 15), (21, 20), (21, 21),
]

# Start of my solution
RUNS = ROW_RUNS + COL_RUNS

BOARD = [[None] * SIZE for _ in xrange(SIZE)]
for (y, x) in GIVENS:
  BOARD[y][x] = True

def to_mask(indexes):
  return reduce(lambda mask, index: mask | 1 << index, indexes, 0)

def to_indexes(mask):
  indexes = []
  bit = 0
  while mask:
    if mask & 1:
      indexes.append(bit)
    mask = mask >> 1
    bit += 1
  return indexes

def sum(l):
  return reduce(lambda a,b: a+b, l, 0)

def all_positions(column):
  runs = RUNS[column]
  spare = SIZE - sum(runs)
  last_positions = [(0, 0, spare)]
  for i, run in enumerate(runs):
    new_positions = []
    for position, offset, spare in last_positions:
      # There's always at least one space before a run, except if it's the first
      for space in range(not not i, spare - len(runs) + i + 2):
        new_positions.append((position | ((2 ** run - 1 | 1) << offset + space),
                              offset + space + run, spare - space))
    last_positions = new_positions
  return [position for position, offset, spare in last_positions]

POSITIONS = [all_positions(i) for i in xrange(SIZE * 2)]

def coords(column, index):
  return (index, column) if column < SIZE else (column - SIZE, index)

def get(column, index):
  x, y = coords(column, index)
  return BOARD[y][x]

def set(column, index, value):
  x, y = coords(column, index)
  BOARD[y][x] = value

def deduce(column):
  runs = RUNS[column]
  defos = to_mask([i for i in range(SIZE) if get(column, i) is True])
  nonos = to_mask([i for i in range(SIZE) if get(column, i) is False])
  positions = POSITIONS[column] = filter(
    lambda position: position & defos == defos and not position & nonos,
    POSITIONS[column]
  )
  assert len(positions), "No possible configurations for " + str(column)
  new_defos = to_indexes(reduce(lambda defos, position: defos & position,
                                positions) & ~defos)
  new_nonos = to_indexes(reduce(lambda nonos, position: nonos & ~position,
                                positions, 2 ** SIZE -1) & ~nonos)
  map(lambda i: set(column, i, True), new_defos)
  map(lambda i: set(column, i, False), new_nonos)
  return len(new_defos) + len(new_nonos)

def draw():
  for row in BOARD:
    print "".join([(u"\u2591" if cell else u"\u2588") * 2
                   for cell in row])

iterations = 1
total = len(GIVENS)
while total < SIZE * SIZE:
  total += sum([deduce(i) for i in xrange(SIZE * 2)])
  iterations += 1
print "Finished after %d iterations." % iterations
draw()