1 # This file is part of NIT ( http://www.nitlanguage.org ).
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
7 # http://www.apache.org/licenses/LICENSE-2.0
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
15 # A trie (or prefix tree) is a datastructure used to perform prefix searches.
17 # The trie uses an arborescent datastructure to perform searches on a prefix.
18 # With this version of the trie, you can link data to nodes so the trie can
19 # be used as a kind of Map by prefix.
22 # # Associate some integers to Map keys
23 # var trie = new Trie[Int]
30 # # Get stored values by key
31 # print trie.has_key("foo")
32 # print trie["foo"] == 1
34 # # Get stored values by prefix
35 # assert trie.has_prefix("fo")
36 # assert not trie.has_prefix("z")
37 # assert trie.find_by_prefix("foo") == [1, 2, 3]
38 # assert trie.find_by_prefix("bar") == [4]
39 # assert trie.find_by_prefix("z").is_empty
45 # Trie data structure for prefix searches
48 # # Associate some integers to Map keys
49 # var trie = new Trie[Int]
55 # assert trie.has_key("foo")
56 # assert trie["foo"] == 1
59 # assert trie.find_by_prefix("") == [1, 3, 2]
60 # assert trie.find_by_prefix("foo") == [1, 2]
61 # assert trie.find_by_prefix("baz").is_empty
68 # Root children are used to store the first letters.
69 private var root
= new TrieNode[E
]
71 redef fun []=(key
, value
) do
72 var children
= root
.children
74 for i
in [0..key
.length
[ do
78 if children
.has_key
(c
) then
81 node
= new TrieNode[E
](c
)
84 children
= node
.children
86 if i
== key
.length
- 1 then
93 # Cache node used between `has_key` and `[]`
94 private var cache
: nullable TrieNode[E
] = null
96 # Cache key used between `has_key` and `[]`
97 private var cache_key
: nullable String = null
100 if cache_key
== key
then return cache
.as(not null).value
.as(not null)
101 var node
= search_node
(key
)
106 redef fun has_key
(key
) do
107 var node
= search_node
(key
)
108 if node
== null then return false
109 var res
= node
.is_leaf
112 cache_key
= key
.as(String)
117 # Find values stored under `prefix`
120 # # Associate some integers to Map keys
121 # var trie = new Trie[Int]
127 # assert trie.find_by_prefix("") == [1, 4, 2, 3]
128 # assert trie.find_by_prefix("foo") == [1, 2, 3]
129 # assert trie.find_by_prefix("bar") == [4]
130 # assert trie.find_by_prefix("baz").is_empty
132 fun find_by_prefix
(prefix
: String): Array[E
] do
137 node
= search_node
(prefix
)
139 if node
== null then return new Array[E
]
140 return node
.collect_values
143 # Find values stored under `prefix`
146 # # Associate some integers to Map keys
147 # var trie = new Trie[Int]
152 # assert trie.has_prefix("")
153 # assert trie.has_prefix("f")
154 # assert not trie.has_prefix("z")
156 fun has_prefix
(prefix
: String): Bool do
157 if prefix
== "" then return true
158 return search_node
(prefix
) != null
161 # Search a node by a key or prefix
163 # Returns `null` if no node matches `string`.
164 private fun search_node
(string
: nullable Object): nullable TrieNode[E
] do
165 if string
== null then return null
167 var children
= root
.children
169 for i
in [0..string
.length
[ do
170 var c
= string
.chars
[i
]
171 if children
.has_key
(c
) then
173 children
= node
.children
183 # TrieNode used to store the Character key of the value
184 private class TrieNode[E
]
186 var value
: nullable E
187 var children
= new HashMap[Char, TrieNode[E
]]
188 var is_leaf
: Bool = false
190 fun collect_values
: Array[E
] do
191 var values
= new Array[E
]
193 var todo
= new List[TrieNode[E
]]
195 while todo
.not_empty
do
196 var node
= todo
.shift
197 var value
= node
.value
198 if value
!= null then values
.add value
199 for child
in node
.children
.values
do