package com.example.myapplication

class Trie {
    //表示Trie数据结构中的一个节点，该节点可以存储其子节点，并且可以标记是否是一个完整单词的结尾
    private data class TrieNode(
        val children: MutableMap<Char, TrieNode> = mutableMapOf(),

        var isEndOfWord: Boolean = false
    )

    private val root = TrieNode()//根节点

    //将一个单词插入到Trie数据结构中。通过遍历单词的每个字符，创建并连接相应的节点，最终在最后一个字符的节点上标记该路径代表一个完整单词。
    fun insert(word: String) {
        var current = root

        for (char in word.lowercase()) {
            current = current.children.getOrPut(char) { TrieNode() }
        }

        current.isEndOfWord = true
    }

    //在Trie数据结构中查找指定的单词是否存在。
    fun search(word: String): Boolean {
        var current = root

        for (char in word.lowercase()) {
            current = current.children[char] ?: return false
        }
        
        return current.isEndOfWord
    }

    //查找以prefix为前缀的所有单词。通过遍历prefix的每个字符，找到相应的节点，然后从该节点开始递归查找所有以该节点为起点的单词。
    fun startsWith(prefix: String): List<String> {
        var current = root

        for (char in prefix.lowercase()) {
            current = current.children[char] ?: return emptyList()
        }
        
        return getAllWordsFromNode(current, prefix)
    }

    //从给定节点开始递归查找所有以该节点为起点的单词。
    private fun getAllWordsFromNode(node: TrieNode, prefix: String): List<String> {
        val words = mutableListOf<String>()

        if (node.isEndOfWord) {
            words.add(prefix)
        }

        for ((char, child) in node.children) {
            words.addAll(getAllWordsFromNode(child, prefix + char))
        }
        
        return words
    }
}
