Create a gist now

Instantly share code, notes, and snippets.

class Trie(object):
"""Simple Trie
>>> t = Trie()
>>> t.add('hi')
>>> t.add('hello')
>>> t.data.keys()
['h']
>>> print t.display(),
h
I
e
l
l
O
>>> t.add('bad')
>>> t.add('ad')
>>> 'hello' in t
True
"""
def __init__(self):
self.data = {}
self.is_word = False
def add(self, word):
cur = self
for c in word:
if c not in cur.data:
cur.data[c] = Trie()
cur = cur.data[c]
cur.is_word = True
def has_children(self, word):
if len(word) == 0:
return bool(self.data)
if word[0] in self.data:
return self.data[word[0]].has_children(word[1:])
return False
def display(self, indent=0):
s = '\n' if indent else ''
for c, t in self.data.items():
s += ' '* indent + (c.upper() if self.data[c].is_word else c.lower()) + t.display(indent+1)
return s
def __getitem__(self, word):
if len(word) == 0:
return self
return self.data[word[0]][word[1:]]
def __contains__(self, word):
if len(word) == 0:
return self.is_word
if word[0] not in self.data:
return False
return word[1:] in self.data[word[0]]
@classmethod
def from_words(cls):
t = Trie()
for word in open('/usr/share/dict/words'):
t.add(word.strip())
return t
def unvisited_neighbors(board, visited, row, col):
"""
>>> unvisited_neighbors(['abc', 'def'], set([(0, 0), (0, 1)]), 0, 0)
[('d', (1, 0)), ('e', (1, 1))]
"""
return [(board[y][x], (y, x))
for y, x in [(row + dy, col + dx)
for dx in [-1, 0, 1]
for dy in [-1, 0, 1]]
if -1 < y < len(board) and -1 < x < len(board[0]) and (y, x) not in visited]
def word_candidates(board):
"""
>>> for cand in word_candidates(['ab']):
... print cand
a
ab
b
ba
"""
for y, row in enumerate(board):
for x, c in enumerate(row):
cands = candidates(board, y, x, visited=set([(y, x)]), candidate=c)
for cand in cands: break
while True:
try:
cand = cands.send((yield cand))
except StopIteration:
break
def candidates(board, row, col, visited, candidate=''):
should_abort = (yield candidate)
if should_abort:
return
for c, (y, x) in unvisited_neighbors(board, visited, row, col):
visited.add((y, x)) #TOMTODO fix this - maontain visited correctly over time
cands = candidates(board, y, x, visited, candidate+c)
for cand in cands: break
while True:
try:
cand = cands.send((yield cand))
except StopIteration:
break
visited.remove((y, x))
import doctest
doctest.testmod()
board = ['abcd',
'efgh',
'ijkl',
'mnop']
t = Trie.from_words()
print t.data['a'].data.keys()
w = word_candidates(board)
candidate = w.next()
while True:
#print 'candidate:', candidate
if candidate in t:
print candidate, 'is a word!'
try:
candidate = w.send(not t.has_children(candidate))
except StopIteration:
break
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment