wordfeud.py

import collections
import functools
import unittest

# import pygame
# class Scrabble:
#     def __init__(self):
#         self.tile_size = 40
#         self.num_tiles = 15
#         self.height = self.tile_size * self.num_tiles
#         self.width = self.tile_size * self.num_tiles
#         self.background = pygame.display.set_mode((800,600))
#         self.background.fill(pygame.Color('black'))
#         self.quiting = False
#         self.draw_board()
#         self.draw_bag()
#         self.main_loop()

#     def draw_board(self):
#         for w in range(0, self.width, self.tile_size):
#             for h in range(0, self.height, self.tile_size):
#                 pygame.draw.rect(
#                     self.background, 
#                     pygame.Color('white'), 
#                     pygame.Rect(w+1, h+1, self.tile_size-2, self.tile_size-2)
#                 )
#         pygame.display.update()

#     def draw_bag(self):
#         w = self.width + self.tile_size
#         for h in range(0, self.tile_size * 7, self.tile_size):
#             pygame.draw.rect(
#                 self.background, 
#                 pygame.Color('red'), 
#                 pygame.Rect(w+1, h+1, self.tile_size-2, self.tile_size-2)
#             )
#         pygame.display.update()

#     def main_loop(self):
#         import time
#         time.sleep(3)

JOKER = "*"

def score(word, scores):
    return sum(scores[l] for l in word)

word_score = functools.partial(
    score,
    scores = {
        'a' : 1, 'b' : 3, 'c' : 5, 'd' : 2, 'e' : 1, 'f' : 4,
        'g' : 3, 'h' : 4, 'i' : 1, 'j' : 4, 'k' : 3, 'l' : 3,
        'm' : 3, 'n' : 1, 'o' : 1, 'p' : 3, 'q' : 10, 'r' : 2,
        's' : 2, 't' : 2, 'u' : 4, 'v' : 4, 'w' : 5, 'x' : 8,
        'y' : 8, 'z' : 4
    }
)

def best_words(trie, letters, fixed):
    return sorted(
        set(trie.get_possible_words(letters, fixed)), 
        key=lambda word_score: (
            -word_score[1],     # Score
            len(word_score[0])  # Word
        )
    )

class Trie:
    __slots__ = ('value', 'children')

    def __init__(self, value=None):
        self.value = value
        self.children = {}

    def __getitem__(self, keys):
        node = self
        for key in keys:
            node = node.children[key]
        if node.value is None:
            raise KeyError('No value for the provided key')
        return node.value

    def __setitem__(self, keys, value):
        cls = type(self)
        node = self
        for key in keys:
            node = node.children.setdefault(key, cls())
        node.value = value

class ScrabbleTrie(Trie):
    def _get_possible_words(self, letters, prefix, fixed):
        if self.value is not None:
            yield prefix, self.value

        char_position = len(prefix)
        if fixed.get(char_position, False):
            for key, node in self.children.items():
                if not fixed[char_position] == key:
                    continue
                yield from node._get_possible_words(letters, prefix + fixed[char_position], fixed)
        else:
            for key, node in self.children.items():
                if not letters.get(key, False):
                    continue
                letters[key] -= 1
                yield from node._get_possible_words(letters, prefix + key, fixed)
                letters[key] += 1

            if letters.get(JOKER, False):
                letters[JOKER] -= 1
                for key, node in self.children.items():
                    yield from node._get_possible_words(letters, prefix + key, fixed)
                letters[JOKER] += 1

    def get_possible_words(self, letters, fixed={}):
        return self._get_possible_words(collections.Counter(letters), '', fixed=fixed)

    def get_best_words(self, letters, fixed={}):
        return best_words(self, letters, fixed=fixed)

# Tests
class TrieTest(unittest.TestCase):
    def setUp(self):
        self.trie = ScrabbleTrie()
        for word in ('tekst', 'test', 'tekens', 'tek', 'tekst', 'teksten'):
            self.trie[word] = word_score(word)

    def test_get_best(self):
        self.assertEqual(
            [('tekst', 10), ('test', 7), ('tek', 6)],
            self.trie.get_best_words(['t', 'k', 'e', 's', 't', 'n'])
        )

    def test_joker(self):
        self.assertEqual(
            [('teksten', 12), ('tekst', 10), ('tekens', 10), ('test', 7), ('tek', 6)],
            self.trie.get_best_words(['t', 'k', 'e', 's', 't', JOKER, 'n'])
        )

    def test_fixed(self):
        self.assertEqual(
            [('tekst', 10), ('test', 7), ('tek', 6)],
            self.trie.get_best_words(['t', 'k', 's', 't'], fixed={1: 'e'})
        )

    def test_fixed_and_joker(self):
        self.assertEqual(
            [('tekst', 10), ('test', 7), ('tek', 6)],
            self.trie.get_best_words(['t', 'k', 's', 't', JOKER], fixed={1: 'e'})
        )

if __name__ == '__main__':
    unittest.main()