LeAf146's blog

1
template <typename T> struct Trie {
2
    struct Node {
3
        unordered_map<T, int> child;
4
        //如果字符集固定，例如只有小写字母，则使用 array<int,26> 通常更快。
5
        int pass = 0, end = 0;
6
    };
7
    vector<Node> tree;
8
    Trie() { tree.emplace_back(); }
9
    template <typename Iterator> void insert(Iterator begin, Iterator end) {
10
        int u = 0;
11
        tree[u].pass++;
12
        for (auto it = begin; it != end; it++) {
13
            const T& c = *it;
14
            auto pos = tree[u].child.find(c);
15
            if (pos == tree[u].child.end()) {
16
                int id = tree.size();
17
                tree[u].child[c] = id;
18
                tree.emplace_back();
19
                u = id;
20
            } else
21
                u = pos->second;
22
            tree[u].pass++;
23
        }
24
        tree[u].end++;
25
    }
26
    template <typename Iterator> int query(Iterator begin, Iterator end) {
27
        int u = 0;
28
        for (auto it = begin; it != end; it++) {
29
            const T& c = *it;
30
            auto pos = tree[u].child.find(c);
31
            if (pos == tree[u].child.end())
32
                return -1;
33
            u = pos->second;
34
        }
35
        return u;
36
    }
37
    // 以上实现必要功能，以下为附加内容
38
    template <typename Container> void insert(const Container& s) { insert(s.begin(), s.end()); }
39
    template <typename Iterator> bool exists(Iterator begin, Iterator end) {
40
        int u = query(begin, end);
41
        return u != -1 && tree[u].end > 0;
42
    }
43
    template <typename Iterator> bool existsPrefix(Iterator begin, Iterator end) {
44
        int u = query(begin, end);
45
        return u != -1 && tree[u].pass > 0;
46
    }
47
    template <typename Iterator> int count(Iterator begin, Iterator end) {
48
        int u = query(begin, end);
49
        if (u == -1)
50
            return 0;
51
        return tree[u].end;
52
    }
53
    template <typename Iterator> bool erase(Iterator begin, Iterator end) {
54
        vector<int> path;
55
        int u = 0;
56
        path.emplace_back(0);
57
        for (auto it = begin; it != end; it++) {
58
            const T& c = *it;
59
            auto pos = tree[u].child.find(c);
60
            if (pos == tree[u].child.end())
61
                return 0;
62
            u = pos->second;
63
            path.emplace_back(u);
64
        }
65
        if (tree[u].end == 0)
66
            return 0;
67
        tree[u].end--;
68
        for (int id : path)
69
            tree[id].pass--;
70
        return 1;
71
    }
72
};