test3x3.py

from creamsteak import creamsteak
import itertools
import numpy as np
import time

def easy_3x3(vals):
    inds = np.argsort(vals)
    ii, jj = np.unravel_index(inds[:3], (3, 3))

    vals.sort()
    total = vals[4:].sum()
    if ii[0] == ii[1] == ii[2] or jj[0] == jj[1] == jj[2]:
        return total + vals[2]
    return total + vals[3]

def hard_3x3(vals):
    return creamsteak(vals.reshape(3, 3))[0]

def do_tests(tests, func=easy_3x3,
             numTests=np.inf, showProg=False, freq=1000,
             name=''):
    print "--------------------------------------------"
    print 'BEGIN TEST:', name
    print 'Maximum # of test cases:', numTests
    print "--------------------------------------------"
    start = time.clock()
    total = 0
    count = 0
    prog = 0
    for vals in tests:
        vals = np.array(vals)
        total += func(vals)
        count += 1
        if count >= numTests:
            break

        if not showProg or numTests == np.inf:
            continue
        newProg = count * freq / numTests
        if newProg > prog:
            prog = newProg
            print 'Completed {} of {} tests ({.2f}%)'.format(count, numTests,
                                                             prog*100.0/freq)
    print "BigBoss =", total
    print "BigCount =", count
    print 'Time taken: {:.2f}s'.format(time.clock()-start)
    print "-----------------END TEST-------------------"
    print

def all_3x3():
    return itertools.product(xrange(1, 10), repeat=9)

def from_file(filename):
    nums = []
    with open(filename) as f:
        for line in f:
            toks = line.split(',')
            if len(toks) > 1:
                for tok in toks:
                    nums.append(int(tok))
            else:
                yield nums
                nums = []
    if len(nums) > 0:
        yield nums

if __name__=='__main__':
    nt = 10**5
    do_tests(all_3x3(), easy_3x3, nt, name='All 3x3 cases w/ shortcut')
    do_tests(all_3x3(), hard_3x3, nt, name='All 3x3 cases using creamsteak()')
    do_tests(from_file("Bunch_o'_matrices.txt"), easy_3x3, name='From file w/ shortcut')
    do_tests(from_file("Bunch_o'_matrices.txt"), hard_3x3, name='From file using creamsteak()')