In [1]:

    

import numpy as np


def move_is_correct(grid,num):
    '''
    @param grid: 6x7 grid containing the current game state
    @param num: column

    returns True if move is allowed on that column
    '''

    #if 0 is in column
    if 0 in grid[:,num]:
    
        #move is allowed
        return True

    else:

        return False

def move_still_possible(S):
    '''
    @param S: 6x7 grid containing the current game state
    returns True if grid contains no 0, therefore no move possible anymore
    '''
    return not(S[S==0].size == 0)


def move(S,p,col_num):
    '''
    @param S: 6x7 grid containing the current game state
    @param p: current player
    @param col_num: column number
    
    sets the player's number on the grid and returns the grid
    '''
    
    #sanity check
    if 0 in S[:,col_num]:    
        y = np.where(S[:,col_num]==0)[0][-1]
        S[y,col_num] = p
        return S , y, col_num
    else:
        return S, None, None
        return 

def move_at_random(S):
    '''
    @param S: 6x7 grid containing the current game state
    moves at random
    '''
    return np.random.randint(0,S.shape[1])


#neat and ugly but the fastest way to search a matrix for a vector is a string find
player1 = '1 1 1 1'
oponent = '2 2 2 2'

def move_was_winning_move(S, p):
    '''
    @param S: 6x7 grid containing the current game state
    @param p: current player
    
    combines all the allowed formations of the grid and string_finds with 
    the currents player vector. Returns true if match.
    '''
    if p == 1:
        match = player1
    else:
        match = oponent 

    l=[]
    #for every possible diag
    for i in range(-2,4):
        l.append(np.diag(S,k = i))
        l.append(np.diag(np.fliplr(S),k=i))
    #left to right
    l.append(S)
    #top to bottom
    l.append(np.rot90(S)) 

    if ''.join(np.array_str(e) for e in l).find(match) > -1:
        return True
    return False



# relate numbers (1, -1, 0) to symbols ('x', 'o', ' ')
symbols = {1:'b', 2:'r', 0:' '}

# print game state matrix using symbols
def print_game_state(S):
    B = np.copy(S).astype(object)
    for n in [1, 2, 0]:
        B[B==n] = symbols[n]
    print B





if __name__ == '__main__':
    # initialize 6x7 connectfour board
    gameState = np.zeros((6,7), dtype=int)

    # initialize player number, move counter
    player = 1
    mvcntr = 1

    # initialize flag that indicates win
    noWinnerYet = True
    

    while move_still_possible(gameState) and noWinnerYet:
        
        while True:
            # get player symbol
            name = symbols[player]
            print '%s moves' % name

            # let player move at random
            col_num = move_at_random(gameState)
            if move_is_correct(gameState, col_num):
                gameState, _ , _ = move(gameState,player,col_num)

                # print current game state
                print_game_state(gameState)

                # evaluate game state
                if move_was_winning_move(gameState, player):
                    print 'player %s wins after %d moves' % (name, mvcntr)
                    noWinnerYet = False
                

                # switch player and increase move counter
                if player == 1:
                    player = 2
                elif player == 2:
                    player = 1

                mvcntr +=  1
                
                break




    if noWinnerYet:
        print 'game ended in a draw'


    

    




b moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' 'b' ' ' ' ' ' ' ' ' ' ']]
r moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' 'b' ' ' 'r' ' ' ' ' ' ']]
b moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' 'b' ' ' ' ' ' ']
 [' ' 'b' ' ' 'r' ' ' ' ' ' ']]
r moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' 'b' ' ' ' ' ' ']
 ['r' 'b' ' ' 'r' ' ' ' ' ' ']]
b moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' 'b' ' ' ' ' ' ']
 ['r' 'b' ' ' 'r' 'b' ' ' ' ']]
r moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' 'r' ' ' 'b' ' ' ' ' ' ']
 ['r' 'b' ' ' 'r' 'b' ' ' ' ']]
b moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' 'r' ' ' 'b' 'b' ' ' ' ']
 ['r' 'b' ' ' 'r' 'b' ' ' ' ']]
r moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' 'r' ' ' 'b' 'b' ' ' ' ']
 ['r' 'b' ' ' 'r' 'b' 'r' ' ']]
b moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' 'r' ' ' 'b' 'b' 'b' ' ']
 ['r' 'b' ' ' 'r' 'b' 'r' ' ']]
r moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' 'r' ' ']
 [' ' 'r' ' ' 'b' 'b' 'b' ' ']
 ['r' 'b' ' ' 'r' 'b' 'r' ' ']]
b moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' 'r' ' ']
 [' ' 'r' ' ' 'b' 'b' 'b' ' ']
 ['r' 'b' 'b' 'r' 'b' 'r' ' ']]
r moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' 'r' ' ' ' ' ' ' 'r' ' ']
 [' ' 'r' ' ' 'b' 'b' 'b' ' ']
 ['r' 'b' 'b' 'r' 'b' 'r' ' ']]
b moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' 'r' ' ' ' ' ' ' 'r' ' ']
 ['b' 'r' ' ' 'b' 'b' 'b' ' ']
 ['r' 'b' 'b' 'r' 'b' 'r' ' ']]
r moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' 'r' ' ' ' ' 'r' 'r' ' ']
 ['b' 'r' ' ' 'b' 'b' 'b' ' ']
 ['r' 'b' 'b' 'r' 'b' 'r' ' ']]
b moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' 'r' ' ' ' ' 'r' 'r' ' ']
 ['b' 'r' ' ' 'b' 'b' 'b' ' ']
 ['r' 'b' 'b' 'r' 'b' 'r' 'b']]
r moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' 'r' ' ' ' ' 'r' 'r' ' ']
 ['b' 'r' ' ' 'b' 'b' 'b' 'r']
 ['r' 'b' 'b' 'r' 'b' 'r' 'b']]
b moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 ['b' 'r' ' ' ' ' 'r' 'r' ' ']
 ['b' 'r' ' ' 'b' 'b' 'b' 'r']
 ['r' 'b' 'b' 'r' 'b' 'r' 'b']]
r moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 ['r' ' ' ' ' ' ' ' ' ' ' ' ']
 ['b' 'r' ' ' ' ' 'r' 'r' ' ']
 ['b' 'r' ' ' 'b' 'b' 'b' 'r']
 ['r' 'b' 'b' 'r' 'b' 'r' 'b']]
b moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 ['r' ' ' ' ' ' ' ' ' ' ' ' ']
 ['b' 'r' ' ' ' ' 'r' 'r' 'b']
 ['b' 'r' ' ' 'b' 'b' 'b' 'r']
 ['r' 'b' 'b' 'r' 'b' 'r' 'b']]
r moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 [' ' ' ' ' ' ' ' ' ' ' ' ' ']
 ['r' ' ' ' ' ' ' ' ' 'r' ' ']
 ['b' 'r' ' ' ' ' 'r' 'r' 'b']
 ['b' 'r' ' ' 'b' 'b' 'b' 'r']
 ['r' 'b' 'b' 'r' 'b' 'r' 'b']]
b moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 ['b' ' ' ' ' ' ' ' ' ' ' ' ']
 ['r' ' ' ' ' ' ' ' ' 'r' ' ']
 ['b' 'r' ' ' ' ' 'r' 'r' 'b']
 ['b' 'r' ' ' 'b' 'b' 'b' 'r']
 ['r' 'b' 'b' 'r' 'b' 'r' 'b']]
r moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 ['b' ' ' ' ' ' ' ' ' ' ' ' ']
 ['r' ' ' ' ' ' ' ' ' 'r' 'r']
 ['b' 'r' ' ' ' ' 'r' 'r' 'b']
 ['b' 'r' ' ' 'b' 'b' 'b' 'r']
 ['r' 'b' 'b' 'r' 'b' 'r' 'b']]
b moves
[[' ' ' ' ' ' ' ' ' ' ' ' ' ']
 ['b' ' ' ' ' ' ' ' ' ' ' ' ']
 ['r' ' ' ' ' ' ' ' ' 'r' 'r']
 ['b' 'r' ' ' ' ' 'r' 'r' 'b']
 ['b' 'r' 'b' 'b' 'b' 'b' 'r']
 ['r' 'b' 'b' 'r' 'b' 'r' 'b']]
player b wins after 23 moves

In [ ]: