Google Code Jam 2013 Round 1B Osmos

'''
Google Code Jam 2013 Round 1B Osmos
@author: Allen

Observations:
- Always add or remove the largest possible mote, and order doesn't matter
- Always have Armin's mote absorb all it can, so I know when I'm done when there are no motes to be absorbed
'''
import os

# Perform an exhaustive search where I calculate the pairs (mote, motes) and
# find the smallest number of ops = add + remove that absorbs all motes.
# This is actually O(N^3) but since N <= 100 I didn't bother to optimize.
def osmos(mote, motes):
    # Makes it simpler and faster to add and remove the largest mote.
    motes.sort()

    # The search space is a 2D array[add][remove]
    # This starting value is where add = remove = 0 (i.e. do nothing)
    space = [[expand(mote, motes)]]

    # The upper bound is where we remove the motes one by one
    upperbound = len(space[0][0][1])

    # I search the smallest number of operations first, which guarantees that
    # the first solution I find is optimal.
    for ops in range(1, upperbound):
        for remove in range(ops):
            add = ops - remove

            if remove == 0:
                # add the largest mote possible = Armin's mote - 1, so 2*mote-1 total
                ref = space[add - 1][0]
                space.append([expand(2 * ref[0] - 1, ref[1])])
            else:
                # remove the largest mote = last mote in a sorted list
                ref = space[add][remove - 1]
                space[add].append((ref[0], ref[1][:-1]))

            # found a solution!
            if(len(space[add][remove][1]) == 0):
                return ops
    return upperbound

# returns the equivalent representation where Armin's mote has absorbed all motes possible
# precondition: motes is sorted
def expand(mote, motes):
    newmote = mote
    bigmotes = []
    for index, mote in enumerate(motes):
        if(mote < newmote):
            newmote += mote
        else:
            # Armin can't absorb the remaining motes
            bigmotes = motes[index:]
            break
    return newmote, bigmotes

# Boilerplate code for I/O
if __name__ == "__main__":
    pname = 'A-example'
    with open(os.path.join(os.path.dirname(__file__), pname+'.in')) as fin, open(os.path.join(os.path.dirname(__file__), pname+'.out'), 'w') as fout:
        T = int(fin.readline())
        for case in range(T):
            mote = int(fin.readline().split()[0])
            result = osmos(mote, [int(i) for i in fin.readline().split()])
            fout.write("Case #{0}: {1}\n".format(case + 1, result))