baduk.py

EMPTY = '.'
BLACK = 'b'
WHITE = 'w'

class Baduk(object):
    def __init__(self, size):
        self.size = size
        self.board = {(i, j): EMPTY
                      for i in range(size)
                      for j in range(size)}
        self.current_player = BLACK
        self.score = {BLACK: 0,
                      WHITE: 0}
        self.process_adjacent_points()

    def __str__(self):
        return '\n'.join(''.join(str(g.board[(i, j)])
                                 for j in range(self.size)) 
                         for i in range(self.size)) + ('\ncurrent player = {}'
                '\nblack score = {}'
                '\nwhite score = {}').format(self.current_player,
                                             self.score[BLACK],
                                             self.score[WHITE])
   
    def process_adjacent_points(self):
        self.adjacent_points = {(0, 0): frozenset([(0, 1), 
                                                   (1, 0)]),
                                (self.size-1, 0): frozenset([(self.size-2, 0), 
                                                             (self.size-1, 1)]),
                                (0, self.size-1): frozenset([(0, self.size-2), 
                                                             (1, self.size-1)]),
                                (self.size-1, self.size-1): frozenset([(self.size-2, self.size-1), 
                                                                       (self.size-1, self.size-2)])}
        for i in range(1, self.size-1):
            self.adjacent_points[(i, 0)] = frozenset([(i-1, 0),
                                                      (i, 1),
                                                      (i+1, 0)])
            self.adjacent_points[(0, i)] = frozenset([(0, i-1),
                                                      (1, i),
                                                      (0, i+1)])
            self.adjacent_points[(i, self.size-1)] = frozenset([(i-1, self.size-1),
                                                                (i, self.size-2),
                                                                (i+1, self.size-1)])
            self.adjacent_points[(self.size-1, i)] = frozenset([(self.size-1, i-1),
                                                                (self.size-2, i),
                                                                (self.size-1, i+1)])
        for i in range(1, self.size-1):
            for j in range(1, self.size-1):
                self.adjacent_points[(i,j)] = frozenset([(i-1, j),
                                                         (i, j-1),
                                                         (i+1, j),
                                                         (i, j+1)])


    @property
    def next_player(self):
        if self.current_player == BLACK:
            return WHITE
        return BLACK

    def alternate_player(self):
        self.current_player = self.next_player

    def liberties(self, point):
        return frozenset(adjacent_point 
                         for adjacent_point in self.adjacent_points[point]
                         if self.board[adjacent_point] == EMPTY)

    def group(self, point):
        color = self.board[point]
        if color == EMPTY:
            return
        g = set()
        explored = set()
        frontier = [point]
        while frontier:
            current_point = frontier.pop()
            if self.board[current_point] == color:
                g.add(current_point)
                for adjacent_point in self.adjacent_points[current_point]:
                    if adjacent_point not in g:
                        frontier.append(adjacent_point)
            explored.add(current_point)
        g.discard(point) # the group doesn't include the original stone
        return g

    def remove(self, points):
        for point in points:
            self.board[point] = EMPTY

    def capture(self, point):
        res = False
        to_remove = set()
        for adjacent_point in self.adjacent_points[point]:
            if (self.board[adjacent_point] == self.next_player) and (len(self.liberties(adjacent_point))-1 == 0):
                current_group = self.group(adjacent_point)
                if not current_group or all(len(self.liberties(stone)) == 0
                                            for stone in current_group):
                    res = True
                    self.score[self.current_player] += len(current_group)+1
                    to_remove.update(current_group)
                    to_remove.add(adjacent_point)
        self.remove(to_remove)
        return res

    def play(self, point):
        if self.board[point] != EMPTY:
            return
       
        if len(self.liberties(point)) == 0:
            if not self.capture(point):
                return # don't allow suicide
        else:
            self.capture(point)

        self.board[point] = self.current_player
        self.alternate_player()


g = Baduk(4)
print(g)
g.play((1, 1))
print(g)
g.play((0, 1))
print(g)
g.play((0, 0))
print(g)
g.play((1, 0))
print(g)
g.play((1, 2))
print(g)
g.play((2, 3))
print(g)
g.play((1, 3))
print(g)
g.play((3, 3))
print(g)
g.play((0, 3))
print(g)
g.play((3, 2))
print(g)
g.play((3, 1))
print(g)
g.play((2, 2))
print(g)
g.play((2, 1))
print(g)