CodeJam Treasure solution

treasure.py

__author__ = "google.com/+robertking"

"""
One of my favorite code jam problems.
solution to the Treasure problem from google code jam 2013:
https://code.google.com/codejam/contest/dashboard?c=2270488#s=p3
"""

from sys import stdin
from collections import Counter, defaultdict


def can_take(i, keys, chests):
    """
    The purpose of this function is to ensure that if we take chest i,
    with the remaining keys and the new keys from chest i,
    is it still possible to reach every key required for the remaining chests..
    If it is still possible, then we can greedily take this chest.
    If there is a certain key that can't be reached then we can't take it.
    """
    idx, Ti, chest_keys = chests[i]
    #firstly.. can we actually open this chest:
    if keys[Ti] == 0:
        return False
    #Now for the search (BFS)... Does the "connected" property hold:
    needed_keys = set(ti for idx2, ti, _ in chests if idx2 != idx)
    keys_by_type = defaultdict(Counter)
    for idx2, Ti2, chest_keys2 in chests:
        if idx2 == idx:
            continue
        keys_by_type[Ti2] += chest_keys2
    fringe = set(keys.keys() + chest_keys.keys())
    if keys[Ti] == 1 and chest_keys[Ti] < 1:
        #we just used Key Ti and it was our last one!
        fringe.remove(Ti)
    #got these keys already (from existing keys or from opened chest)..
    needed_keys -= fringe
    while needed_keys:
        new_fringe = set()
        for parent in fringe:
            for child in keys_by_type[parent]:
                if child in needed_keys:
                    needed_keys.remove(child)
                    new_fringe.add(child)
        if new_fringe:
            fringe = new_fringe
        else:
            return False if needed_keys else True
    return True


def solve(keys, chests, N):
    if N == 0:
        return []
    for i in range(len(chests)):
        if can_take(i, keys, chests):
            idx, Ti, chest_keys = chests.pop(i)
            keys -= Counter([Ti])
            keys += chest_keys
            return [idx] + solve(keys, chests, N-1)


nums = (map(int, line.split()) for line in stdin)
for case in range(1, next(nums)[0] + 1):
    K, N = next(nums)
    keys = Counter(next(nums))
    types = Counter()
    chests = []
    for i in range(1, N + 1):
        x = next(nums)
        Ti, Ki, chest_keys = x[0], x[1], Counter(x[2:])
        chests.append((i, Ti, chest_keys))
        types[Ti] += 1
    #clean keys
    for k in keys.keys():
        if k not in types:
            del keys[k]
    #clean chests
    for chest in chests:
        for ti in chest[2].keys():
            if ti not in types:
                del chest[2][ti]
    #what keys do we have available.
    available = keys.copy()
    for chest in chests:
        available += chest[2]
    #do we have enough keys for all the chests:
    if types - available:
        ans = None
    else:
        ans = solve(keys, chests, N)
    print "Case #%d: %s" % (case, "IMPOSSIBLE" if ans is None else " ".join(str(i) for i in ans))