1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

//! トライ木

use std::fmt::Debug;

// 定数
const ORIGIN: char = 'a'; // 基準となる文字
const ORIGIN_ID: usize = ORIGIN as usize; // 基準となる文字のID
const KINDS: usize = 26; // 文字の種類数
type NodePointer<T> = Option<Box<TrieNode<T>>>;

/// 何番目の文字かを判定する
fn ord(c: char) -> usize {
    let num = c as usize;
    num - ORIGIN_ID
}

/// i番目の文字を返す
fn chr(i: usize) -> char {
    (ORIGIN_ID + i) as u8 as char
}

/// # TrieNode
/// - トライ木のノード
#[derive(Debug, Clone)]
struct TrieNode<T> {
    data: Option<T>,
    children: Vec<NodePointer<T>>,
}

impl<T> TrieNode<T>
where
    T: Clone,
{
    pub fn new(data: Option<T>) -> Self {
        Self {
            data,
            children: vec![NodePointer::None; KINDS],
        }
    }
}

/// # Trie
/// - トライ木の実装
#[derive(Debug)]
pub struct Trie<T> {
    size: usize,
    root: NodePointer<T>,
}

impl<T> Trie<T>
where
    T: Clone + Debug,
{
    pub fn len(&self) -> usize {
        self.size
    }

    pub fn is_empty(&self) -> bool {
        self.size == 0
    }

    pub fn insert(&mut self, key: &str, data: T) -> Option<T> {
        let res = self.get_or_insert_mut(key).replace(data);
        if res.is_none() {
            self.size += 1;
        }
        res
    }

    pub fn get(&self, key: &str) -> Option<&T> {
        let mut node = &self.root;
        for c in key.chars().map(ord) {
            node = &node.as_ref()?.children[c];
        }
        node.as_deref()?.data.as_ref()
    }

    pub fn get_mut(&mut self, key: &str) -> Option<&mut T> {
        let mut node = &mut self.root;
        for c in key.chars().map(ord) {
            node = node.as_mut()?.children.get_mut(c).unwrap();
        }
        node.as_deref_mut()?.data.as_mut()
    }

    pub fn get_or_insert_mut(&mut self, key: &str) -> &mut Option<T> {
        let mut node = &mut self.root;
        for c in key.chars().map(ord).chain(KINDS..=KINDS) {
            // データの挿入
            if c == KINDS {
                if node.as_ref().is_none() {
                    *node = Some(Box::new(TrieNode::new(None)));
                }
                break;
            }
            if node.as_ref().is_none() {
                *node = Some(Box::new(TrieNode::new(None)));
            }
            node = node.as_mut().unwrap().children.get_mut(c).unwrap();
        }
        &mut node.as_deref_mut().unwrap().data
    }

    pub fn traverse(&self) -> Vec<(String, &T)> {
        let mut res = vec![];
        let mut cur = String::new();
        traverse_inner(&self.root, &mut cur, &mut res);
        res
    }
}

impl<T: Clone> Default for Trie<T> {
    fn default() -> Self {
        Trie {
            size: 0,
            root: Some(Box::new(TrieNode {
                data: None,
                children: vec![NodePointer::None; KINDS],
            })),
        }
    }
}

/// trieを順に探索する
fn traverse_inner<'a, T>(
    node: &'a NodePointer<T>,
    cur: &mut String,
    list: &mut Vec<(String, &'a T)>,
) {
    if let Some(value) = node.as_ref().unwrap().data.as_ref() {
        let key = cur.clone();
        list.push((key, value));
    }
    if let Some(node) = node.as_deref() {
        for (i, child) in node.children.iter().enumerate() {
            if child.as_ref().is_some() {
                cur.push(chr(i));
                traverse_inner(child, cur, list);
                cur.pop();
            }
        }
    }
}