字典树的实现

字典树的实现_牛客题霸_牛客网

问题简述

字典树又称为前缀树或者Trie树，是处理字符串常用的数据结构。

假设组成所有单词的字符仅是‘a’～‘z’，请实现字典树的结构，并包含以下四个主要的功能。

1. void insert(String word)：添加word，可重复添加；
2. void delete(String word)：删除word，如果word添加过多次，仅删除一次；
3. boolean search(String word)：查询word是否在字典树中出现过(完整的出现过，前缀式不算)；
4. int prefixNumber(String pre)：返回以字符串pre作为前缀的单词数量。

现在给定一个m，表示有m次操作，每次操作都为以上四种操作之一。每次操作会给定一个整数op和一个字符串word，op代表一个操作码，如果op为1，则代表添加word，op为2则代表删除word，op为3则代表查询word是否在字典树中，op为4代表返回以word为前缀的单词数量（数据保证不会删除不存在的word）。

对于每次操作，如果op为3时，如果word在字典树中，请输出“YES”，否则输出“NO”；如果op为4时，请输出返回以word为前缀的单词数量，其它情况不输出。

进阶：所有操作的时间复杂度都满足 O(n)

思路

• 本质上是一个 N 叉树，根据需求，为每个节点添加一些属性以达到某些操作的要求；

  class TrieNode:
      def __init__(self):
          self.children = [None] * 26  # 子节点（词表大小为 26）
          self.is_end = False          # 该节点是否为单词
          self.pre_cnt = 0             # 记录该路径的作为前缀的次数

• 相同实现 Python 实现会超时；Java 不会；

Python（超时）

class TrieNode:
    def __init__(self):
        self.is_end = False
        self.children = [None] * 26
        self.pre_cnt = 0

class Trie:
    
    root = TrieNode()
    
    def insert(self, word):
        cur = self.root
        for c in word:
            idx = ord(c) - ord('a')
            if not cur.children[idx]:
                cur.children[idx] = TrieNode()
            cur = cur.children[idx]
            cur.pre_cnt += 1
        cur.is_end = True
    
    def delete(self, word):
        # 只有存在的情况再进行删除
        if not self.search(word): return
            
        cur = self.root
        for c in word:
            idx = ord(c) - ord('a')
            cur = cur.children[idx]
            cur.pre_cnt -= 1
        
        if cur.pre_cnt == 0:
            cur.is_end = False
    
    def search(self, word):
        cur = self.root
        for c in word:
            idx = ord(c) - ord('a')
            if not cur.children[idx]:
                return False
            cur = cur.children[idx]
            
        return cur.is_end
    
    def prefixNumber(self, pre):
        cur = self.root
        for c in pre:
            idx = ord(c) - ord('a')
            if not cur.children[idx]:
                return 0
            cur = cur.children[idx]
            if cur.pre_cnt == 0:
                return 0
        
        return cur.pre_cnt

    
class Solution:
    def trieU(self , operators: List[List[str]]) -> List[str]:
        
        trie = Trie()
        ret = []
        for i, w in operators:
            if i == '1':
                trie.insert(w)
            elif i == '2':
                trie.delete(w)
            elif i == '3':
                r = 'YES' if trie.search(w) else 'NO'
                ret.append(r)
            else:
                r = str(trie.prefixNumber(w))
                ret.append(r)
        
        return ret

Java

import java.util.*;

class TrieNode {
    TrieNode[] children;  // 子节点
    boolean is_end;
    int pre_cnt;
    
    TrieNode() {
        children = new TrieNode[26];
        pre_cnt = 0;
        is_end = false;
    }
}

class Trie {
    TrieNode root = new TrieNode();

    Trie() { }

    void insert(String word) {
        TrieNode node = root;
        for (char c : word.toCharArray()) {
            if (node.children[c - 'a'] == null) {
                node.children[c - 'a'] = new TrieNode();
            }
            node = node.children[c - 'a'];
            node.pre_cnt++;
        }
        node.is_end = true;
    }

    void delete(String word) {
        // 只有存在的情况再进行删除（虽然没有这句也能 AC）
        if (!search(word)) return;

        TrieNode node = root;
        for (char c : word.toCharArray()) {
            node = node.children[c - 'a'];
            node.pre_cnt--;
        }
        if (node.pre_cnt == 0) {
            node.is_end = false;
        }
    }

    boolean search(String word) {
        TrieNode node = root;
        for (char c : word.toCharArray()) {
            if (node.children[c - 'a'] == null) {
                return false;
            }
            node = node.children[c - 'a'];
        }

        return node.is_end;
    }

    int prefixNumber(String pre) {
        TrieNode node = root;
        for (char c : pre.toCharArray()) {
            if (node.children[c - 'a'] == null) {
                return 0;
            }
            node = node.children[c - 'a'];
        }

        return node.pre_cnt;
    }
}

public class Solution {
    /**
     * @param operators string字符串二维数组 the ops
     * @return string字符串一维数组
     */
    public String[] trieU(String[][] operators) {

        ArrayList<String> ret = new ArrayList<>();
        Trie trie = new Trie();

        for (String[] opera : operators) {
            if (opera[0].equals("1")) {
                trie.insert(opera[1]);
            } else if (opera[0].equals("2")) {
                trie.delete(opera[1]);
            } else if (opera[0].equals("3")) {
                ret.add(trie.search(opera[1]) ? "YES" : "NO");
            } else if (opera[0].equals("4")) {
                ret.add(String.valueOf(trie.prefixNumber(opera[1])));
            }
        }
        
        String[] ans = new String[ret.size()];
        ret.toArray(ans);
        return ans;
    }
}