Возвращает строки, соответствующие префиксу в Trie.

Question

Мне удалось создать Trie, и теперь я хочу вернуть строки, соответствующие префиксу из Trie, но у меня возникают проблемы при написании функции поиска.

Например, если у меня есть префикс "aa", я хочу иметь строки "aa" и "aac" в качестве вывода.

class Node:
    def __init__(self):
        self.children = [None] * 26
        self.end = False
        self.value = ""

class Trie:
    def __init__(self):
        self.root = Node()

    def add_word(self, key):
        word_length = len(key)
        current = self.root
        for i in range(word_length):
            position = self.ord_char(key[i])

            if current.children[position] is None:
                current.children[position] = Node()
            current = current.children[position]
            current.value = key[i]
        current.end = True

    def ord_char(self,key):
        ord_rep = ord(key) - ord('a')
        return ord_rep

    def prefix_search(self, prefix):
        lst = []
        current = self.root
        prefix_length = len(prefix)
        for c in range(prefix_length):
            c_position = self.ord_char(prefix[c])
            current = current.children[c_position]
            lst.append(current.value)
        #doesnt seem like I'm doing it right

if __name__ == "__main__":
    trie = Trie()
    trie.add_word("aa")
    trie.add_word("aac")
    trie.add_word("b")
    trie.prefix_search("aa")

Я думал о том, чтобы просто объединить алфавиты, чтобы сформировать окончательную строку с помощью функции поиска, но я просто не мог придумать лучшего способа сделать это.

Answer 1

До сих пор значение lst является просто префиксом, разделенным на отдельные буквы, но теперь вам нужно обработать каждый узел, который вы найдете в атрибуте children, который не является None, чтобы найти все узлы, которые имеют end установлен на True. Каждый раз, когда вы находите такой узел, у вас есть полное слово. Любой узел может снова иметь несколько дочерних узлов, разветвляясь на еще больше слов для вывода.

Вы можете использовать стек для отслеживания всех узлов, которые вам нужно обработать для построения списка, вместе с их префиксом до сих пор. Добавьте дочерние узлы в стек с префиксом для этого узла и обрабатывайте эти узлы один за другим (добавляя больше дочерних узлов в стек при этом).

Обратите внимание, что для начала вам не нужно составлять список префиксных символов, у вас уже есть этот префикс в качестве переменной. Чтобы добраться до вашей начальной точки, проще просто перебрать сам префикс:

def prefix_search(self, prefix):
    current = self.root
    # get to the starting point
    for c in prefix:
        current = current.children[self.ord_char(c)]
        if current is None:
            # prefix doesn't exist, abort with an empty list
            return []

    found = []
    stack = [(current, prefix)]
    while stack:
        current, prefix = stack.pop()

        if current.end:
            # this is a complete word, named by prefix
            found.append(prefix)

        # add the children to the stack, each with their letter added to the
        # prefix value.
        for child in current.children:
            if child is None:
                continue
            stack.append((child, prefix + child.value))

    return found

Для вашего заданного примера дерева и префикса стек начинается с узла в 'aa'. Первая итерация while stack: удаляет этот узел из стека, и поскольку для этого узла end установлено значение true, 'aa' добавляется к found. Узел имеет только один не None дочерний узел для c, поэтому узел добавляется в стек с 'aac'.

Затем цикл while повторяется, находит, что один элемент в стеке, видит, что end установлен так, что 'aac' добавляется к found, и больше не обнаруживаются дочерние узлы. Стек остается пустым, и цикл while заканчивается.

Демо-версия:

>>> trie = Trie()
>>> trie.add_word("aa")
>>> trie.add_word("aac")
>>> trie.add_word("b")
>>> trie.prefix_search("aa")
['aa', 'aac']
>>> trie.prefix_search("b")
['b']
>>> trie.add_word('abracadabra')
>>> trie.add_word('abbreviation')
>>> trie.add_word('abbreviated')
>>> trie.add_word('abbrasive')
>>> trie.prefix_search("ab")
['abracadabra', 'abbreviation', 'abbreviated', 'abbrasive']
>>> trie.prefix_search("abr")
['abracadabra']
>>> trie.prefix_search("abb")
['abbreviation', 'abbreviated', 'abbrasive']
>>> trie.prefix_search("abbra")
['abbrasive']

Answer 2

Как насчет .startswith(), мне кажется, простой способ реализовать ваш поиск.