leetcode208：实现Trie

题目链接

leetcode

题目描述

Trie 或者说 前缀树 是一种树形数据结构，用于高效地存储和检索字符串数据集中的键。这一数据结构有相当多的应用情景，例如自动补完和拼写检查。

请你实现 Trie 类：

• Trie() 初始化前缀树对象。
• void insert(String word) 向前缀树中插入字符串 word 。
• boolean search(String word) 如果字符串 word 在前缀树中，返回 true（即，在检索之前已经插入）；否则，返回 false 。
• boolean startsWith(String prefix) 如果之前已经插入的字符串 word 的前缀之一为 prefix ，返回 true ；否则，返回 false 。

C++ 代码

#include <iostream>
#include <vector>
#include <queue>
using namespace std;

/*
    Trie (前缀树)

    1.定义树中结点的结构
        每个结点包含 26 个子节点的指针，bool 变量表示是否为一个单词的结尾。
    2.插入时沿着根结点向下寻找
        若不存在该路径，则创建结点来生成路径。
    3.查询或查询前缀时，沿着路径查找。

    时间复杂度：O(s)
        单次操作的时间复杂度和字符串长度一致。
        插入操作：O(s)，s 是待插入的字符串长度。
        查询操作：O(s)，s 是待查询的字符串长度。
        查询前缀操作：O(s)，s 是待查询的前缀长度。
    空间复杂度：O(nK)
        n 是结点个数，K 是每个结点的分枝数（字符集的大小）。
*/
class Trie {
public:
    Trie() {
        root = new TrieNode();
    }
    
    // 析构函数（工程中需要）
    ~Trie() {
        queue<TrieNode*> q;
        q.push(root);
        while (!q.empty()) {
            TrieNode* node = q.front(); q.pop();
            // for (int i = 0; i < 26; i ++) {
            for (int i = 0; i < TrieNode::N; i ++) {
                if (node->child[i] != nullptr) {
                    q.push(node->child[i]);
                }
            }

            delete node;
        }
    }
    
    void insert(string word) {
        TrieNode* p = root;
        for (char c : word) {
            int index = c - 'a';
            if (p->child[index] == nullptr) {
                p->child[index] = new TrieNode();
            }
            p = p->child[index];
        }
        p->isEnd = true;
    }
    
    bool search(string word) {
        TrieNode* p = root;
        for (char c : word) {
            int index = c - 'a';
            if (p->child[index] == nullptr) {
                return false;
            }
            p = p->child[index];
        }
        return p->isEnd;
    }
    
    bool startsWith(string prefix) {
        TrieNode* p = root;
        for (char c : prefix) {
            int index = c - 'a';
            if (p->child[index] == nullptr) {
                return false;
            }
            p = p->child[index];
        }
        return true;
    }

private:
    struct TrieNode {
        static const int N = 26;
        TrieNode* child[N];
        bool isEnd;

        TrieNode() {
            for (int i = 0; i < N; i ++) {
                child[i] = nullptr;
            }
            isEnd = false;
        }
    };

    TrieNode* root;
};

int main() {
    Trie* trie = new Trie();
    trie->insert("apple");
    cout << trie->search("apple") << endl;   // 返回 true
    cout << trie->search("app") << endl;     // 返回 false
    cout << trie->startsWith("app") << endl; // 返回 true
    trie->insert("app");
    cout << trie->search("app") << endl;     // 返回 true

    return 0;
}

Golang 代码

// 208. 实现 Trie (前缀树)

package main

import (
	"fmt"
)

// TrieNode 定义 Trie 的节点
type TrieNode struct {
	children map[rune]*TrieNode
	isEnd    bool
}

// Trie 定义 Trie 数据结构
type Trie struct {
	root *TrieNode
}

// Constructor 初始化 Trie
func Constructor() Trie {
	return Trie{
		root: &TrieNode{
			children: make(map[rune]*TrieNode),
		},
	}
}

// Insert 插入一个单词到 Trie 中
func (this *Trie) Insert(word string) {
	node := this.root
	for _, ch := range word {
		if _, exists := node.children[ch]; !exists {
			node.children[ch] = &TrieNode{
				children: make(map[rune]*TrieNode),
			}
		}
		node = node.children[ch]
	}
	node.isEnd = true
}

// Search 搜索一个单词是否存在于 Trie 中
func (this *Trie) Search(word string) bool {
	node := this.root
	for _, ch := range word {
		if _, exists := node.children[ch]; !exists {
			return false
		}
		node = node.children[ch]
	}
	return node.isEnd
}

// StartsWith 搜索是否存在以 prefix 开头的单词
func (this *Trie) StartsWith(prefix string) bool {
	node := this.root
	for _, ch := range prefix {
		if _, exists := node.children[ch]; !exists {
			return false
		}
		node = node.children[ch]
	}
	return true
}

func main() {
	testCases := []struct {
		actions  []string
		params   [][]string
		expected []interface{}
	}{
		{
			actions:  []string{"Trie", "insert", "search", "search", "startsWith", "insert", "search"},
			params:   [][]string{{}, {"apple"}, {"apple"}, {"app"}, {"app"}, {"app"}, {"app"}},
			expected: []interface{}{nil, nil, true, false, true, nil, true},
		},
		{
			actions:  []string{"Trie", "insert", "search", "startsWith"},
			params:   [][]string{{}, {"hello"}, {"hello"}, {"he"}},
			expected: []interface{}{nil, nil, true, true},
		},
	}

	for i, tc := range testCases {
		trie := Constructor()
		output := []interface{}{nil} // 初始化输出列表

		for j := 1; j < len(tc.actions); j++ {
			action := tc.actions[j]
			param := tc.params[j]

			switch action {
			case "insert":
				trie.Insert(param[0])
				output = append(output, nil)
			case "search":
				result := trie.Search(param[0])
				output = append(output, result)
			case "startsWith":
				result := trie.StartsWith(param[0])
				output = append(output, result)
			}
		}

		fmt.Printf("Test Case %d, Input: \n", i+1)
        fmt.Printf("actions = %v\n", tc.actions)
        fmt.Printf("params = %v\n", tc.params)
        outputStr := fmt.Sprint(output)
        expectedStr := fmt.Sprint(tc.expected)

        if outputStr == expectedStr {
			fmt.Printf("Test Case %d, Output: %v, PASS\n", i+1, outputStr)
        } else {
            fmt.Printf("Test Case %d, Output: %v, FAIL (Expected: %d)\n", i+1, outputStr, expectedStr)		
        }
	}
}