Otra de exploración de árboles de palabras

El lunes pasado se cayó Internet en el curro. Mientras yo estaba con mis pajas mentales, mis queridos compañeros nerds se empezaron a aglomerar en torno a un ordenador. Para no desentonar, me acerqué yo también y pregunté que qué regalaban.

Resultó que estaban jugando con un programita de consola (no me refiero a xbox ni ps3, sino a terminal de texto) que muestra una parrilla con letras y nosotros tenemos que introducir todas las palabras inglesas que se puedan formar con las letras siguiendo un camino de casilla en casilla.

Acordándome del mítico post sobre El Melenas de Cifras y Letras y de cómo le quité la gracia al juego del facebook de hacer anagramas, decidí que iba petarlo en cuanto volviese Internet y pudiese descargarme la lista de palabras inglesas...

Total, que como buen procrastinador, me puse a hacer mi script en Python, que ahora pasaré a explicar:

from itertools import product
import sys

# not much shorter, but much funnier to write than
# [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)]
# at first I used range(-1,2), but (-1,0,1) is cleaner and shorter
neighbours = [(i,j) for i,j in product(*[(-1,0,1)]*2) if i or j]

allow_repeat = True

def make_tree(words_list):
        base = {}
        for word in words_list:
                d = base
                for c in word:
                        if c not in d:
                                nd = {}
                                d[c] = nd
                                d = nd
                        else:
                                d = d[c]
                d[None] = None # complete word marker
        return base

def parse_matrix(text):
      return map(str.upper, filter(None, (filter(str.isalpha, line) for line in text)))

def is_valid_and_prefix(d, word):
        for c in word:
                if c not in d:
                        return False, False
                else:
                        d = d[c]
        valid = None in d
        prefix = len(d) > valid # more keys than None
        return (valid, prefix)

def successors(word, path, matrix):
        i,j = path[-1]
        for di, dj in neighbours:
                ni, nj = i+di, j+dj
                if (ni >= 0 and nj >= 0 and
                    ni < len(matrix) and nj < len(matrix[ni]) and
                    (allow_repeat or (ni,nj) not in path)):
                        successor = word + matrix[ni][nj], path + [(ni,nj)]
                        yield successor

def find_words(d, path, word, matrix):
        valid, prefix = is_valid_and_prefix(d, word)
        if valid:
                yield word
        if prefix:
                for nw, np in successors(word, path, matrix):
                        for x in find_words(d, np, nw, matrix):
                                yield x

if __name__ == '__main__':
        wordlist = sys.argv[1]
        tree = make_tree(x.upper().strip() for x in open(wordlist))
        matrix = parse_matrix(sys.stdin.readlines())
        words = set()
        for i,j in product(range(len(matrix)), range(len(matrix[0]))):
                c = matrix[i][j]
                words |= set(find_words(tree, [(i,j)], c, matrix))
        for w in sorted(words, key=len):
                print w

Para usarlo, recibe como primer y único argumento el nombre del fichero con el corpus de palabras válidas y lee por la entrada estándar la "sopa de letras" (de la cuál descartará cualquier caracter que no sea alfabético, así que puede estar formateada como os dé la gana).

¿Cómo funciona el jodío? Pues muy fácil, esta vez en lugar de construir un hash de las palabras que ignorase el orden de aparición de las letras, nuestra estructura de consulta es un arbolito (como bien indica el nombre de la función make_tree).

Este arbolito es básicamente un puñado de diccionarios (no en el sentido de RAE, sino tablas asociativas) anidados, en el cuál la clave es una letra y el valor, pues otro diccionaro. Hay un caso especial en el que algunos diccionarios tendrán una entrada None:None.

Para saber si una palabra está en el diccionario (esta vez sí, en el sentido de RAE, Collins, Oxford, etc.), tenemos que ver si podemos llegar a un nodo que tenga una clave None pegando saltitos por las ramas que tengan como clave cada letra de nuestra palabra de forma consecutiva. Esto es lo que hace la función is_valid_and_prefix, que además nos dice si la palabra es un prefijo de otras.

A partir de aquí es pan comido, sólo tenemos que hacer el típico backtracking por las casillas de la rejilla, podando cuando lleguemos a palabras que no son prefijo de ninguna y devolviendo aquellas que sean palabras completas. Los generadores de Python ayudan mucho para escribir este código de forma elegante.

Posibles optimizaciones:

• Reemplazar los diccionarios por alguna otra estructura más compacta en memoria, ya que buscar linealmente entre un máximo de 28 caracteres seguramente sea más eficiente que un hash lookup.
• Pasar sólo el sub-árbol con los sufijos de la palabra que estamos explorando en cada nivel de la búsqueda... No es que vaya a ser una gran diferencia, pero si ya estamos seguros de que vamos por un prefijo válido ¿para qué comprobarlo otra vez?

Lo más gracioso es la cantidad de palabras que se pueden formar con una matriz pequeñita si permitimos repetir casillas. Para probar, usé esta matriz de entrada (sí, ya sé que tiene algunas palabras muy obvias):

DONE

OIRT

PASS

¿Y a que no sabéis cuántas palabras había escondidas? 848.