gistfile1.pyw

#!/usr/bin/env python

'''
Call the function who_wins and pass an array, e.g.:
    who_wins([[1,0,1],
              [1,2,2],
              [1,2,2]])
and it will return 1 or 2, corresponding to a tic-tac-toe
win by either player 1 or player 2 (or, in fact, a triplet
of any repeating character).

It will return 0 for no winner.
'''

def who_wins(tic_tac_toe):
    return reduce(lambda x,y: x or y,[check(tic_tac_toe) for check in checks])

def is_winning_seq(seq):
    ''' If all elements of seq[] are the same, return that element.
    Otherwise, return 0. '''
    s = set(seq)
    return (len(s) == 1 or 0) and s.pop()

def check_rows(array):
    for row in array:
        is_win = is_winning_seq(row)
        if is_win:
            break
    return is_win

def check_cols(array):
    for col in range(len(array)):
        is_win = is_winning_seq([row[col] for row in array])
        if is_win:
            break
    return is_win

def check_diag(array,rows=[],cols=[]):
    '''rows/cols: [0, 1, 2] / [2, 1, 0] for
    a top-right to borrom-left diagonal'''
    seq = []
    for row,col in zip(rows, cols):
        seq.append(array[row][col])
    return is_winning_seq(seq)

def check_diags(array):
    rows = range(len(array))
    return check_diag(array, rows, rows[::-1]) \
        or check_diag(array, rows, rows)

checks = [check_cols, check_rows, check_diags]

### Tests

boards = [
    ([ [1,0,1],
       [1,2,2],
       [1,2,2] ], 1, 'wins with a column'),

    ([ [1,0,1],
       [1,2,1],
       [2,2,2] ], 2, 'wins with a row'),

    ([ [2,0,1],
       [1,2,1],
       [1,2,2] ], 2, 'wins with a diagonal'),

    ([ [2,0,1],
       [1,1,2],
       [1,2,2] ], 1, 'wins with a diagonal'),

    ([ [1,0,1],
       [2,1,2],
       [0,2,2] ], 0, 'nobody wins')
]

def ttt_test():
    for board, expected_result, comment in boards:
        result = who_wins(board)
        assert result  == expected_result, 'unexpected result: %s' % result
        print '%s %s' % (expected_result, comment)
    print 'Passed all tests!'

if __name__ == '__main__':
    ttt_test()