Merge pull request #756 from lucidfrontier45/trie-improve

Improvements of Trie Structure

Author: oxinabox

docs/src/trie.md

@@ -31,3 +31,10 @@ given string, use:
 ```julia
 seen_prefix(t::Trie, str) = any(v -> v.is_key, partial_path(t, str))
 ```
+
+`find_prefixes` can be used to find all keys which are prefixes of the given string.
+
+```julia
+t = Trie(["A", "ABC", "ABCD", "BCE"])
+find_prefixes(t, "ABCDE") # "A", "ABC", "ABCD"
+```

src/DataStructures.jl

@@ -28,7 +28,7 @@ module DataStructures
     export heapify!, heapify, heappop!, heappush!, isheap
     export BinaryMinMaxHeap, popmin!, popmax!, popall!
 
-    export Trie, subtrie, keys_with_prefix, partial_path
+    export Trie, subtrie, keys_with_prefix, partial_path, find_prefixes
 
     export LinkedList, Nil, Cons, nil, cons, head, tail, list, filter, cat,
            reverse

src/trie.jl

@@ -111,14 +111,43 @@ end
 # since the root of the trie corresponds to a length 0 prefix of str.
 function Base.iterate(it::TrieIterator, (t, i) = (it.t, 0))
     if i == 0
-        return it.t, (it.t, 1)
-    elseif i == length(it.str) + 1 || !(it.str[i] in keys(t.children))
+        return it.t, (it.t, firstindex(it.str))
+    elseif i > lastindex(it.str) || !(it.str[i] in keys(t.children))
         return nothing
     else
         t = t.children[it.str[i]]
-        return (t, (t, i + 1))
+        return (t, (t, nextind(it.str, i)))
     end
 end
 
 partial_path(t::Trie, str::AbstractString) = TrieIterator(t, str)
 Base.IteratorSize(::Type{TrieIterator}) = Base.SizeUnknown()
+
+"""
+    find_prefixes(t::Trie, str::AbstractString)
+
+Find all keys from the `Trie` that are prefix of the given string
+
+# Examples
+```julia-repl
+julia> t = Trie(["A", "ABC", "ABCD", "BCE"])
+
+julia> find_prefixes(t, "ABCDE")
+3-element Vector{AbstractString}:
+ "A"
+ "ABC"
+ "ABCD"
+```
+"""
+function find_prefixes(t::Trie, str::AbstractString)
+    prefixes = AbstractString[]
+    it = partial_path(t, str)
+    idx = 0
+    for t in it
+        if t.is_key
+            push!(prefixes, str[firstindex(str):idx])
+        end
+        idx = nextind(str, idx)
+    end
+    return prefixes
+end

test/test_trie.jl

@@ -40,4 +40,30 @@
         @test collect(partial_path(t, "ro")) == [t0, t1, t2]
         @test collect(partial_path(t, "roa")) == [t0, t1, t2]
     end
+
+    @testset "partial_path iterator non-ascii" begin
+        t = Trie(["東京都"])
+        t0 = t
+        t1 = t0.children['東']
+        t2 = t1.children['京']
+        t3 = t2.children['都']
+        @test collect(partial_path(t, "西")) == [t0]
+        @test collect(partial_path(t, "東京都")) == [t0, t1, t2, t3]
+        @test collect(partial_path(t, "東京都渋谷区")) == [t0, t1, t2, t3]
+        @test collect(partial_path(t, "東京")) == [t0, t1, t2]
+        @test collect(partial_path(t, "東京スカイツリー")) == [t0, t1, t2]
+    end
+    
+    @testset "find_prefixes" begin
+        t = Trie(["A", "ABC", "ABD", "BCD"])
+        prefixes = find_prefixes(t, "ABCDE")
+        @test prefixes == ["A", "ABC"]
+    end
+
+    @testset "find_prefixes non-ascii" begin
+        t = Trie(["東京都", "東京都渋谷区", "東京都新宿区"])
+        prefixes = find_prefixes(t, "東京都渋谷区東")
+        @test prefixes == ["東京都", "東京都渋谷区"]
+    end
+
 end # @testset Trie