Merge pull request #759 from jonas-schulze/vector-trie

Allow non-string indices for Trie

Author: oxinabox

src/trie.jl

@@ -1,60 +1,56 @@
-mutable struct Trie{T}
-    value::T
-    children::Dict{Char,Trie{T}}
+mutable struct Trie{K,V}
+    value::V
+    children::Dict{K,Trie{K,V}}
     is_key::Bool
 
-    function Trie{T}() where T
-        self = new{T}()
-        self.children = Dict{Char,Trie{T}}()
+    function Trie{K,V}() where {K,V}
+        self = new{K,V}()
+        self.children = Dict{K,Trie{K,V}}()
         self.is_key = false
         return self
     end
 
-    function Trie{T}(ks, vs) where T
-        t = Trie{T}()
-        for (k, v) in zip(ks, vs)
-            t[k] = v
-        end
-        return t
+    function Trie{K,V}(ks, vs) where {K,V}
+        return Trie{K,V}(zip(ks, vs))
     end
 
-    function Trie{T}(kv) where T
-        t = Trie{T}()
+    function Trie{K,V}(kv) where {K,V}
+        t = Trie{K,V}()
         for (k,v) in kv
             t[k] = v
         end
         return t
     end
 end
 
-Trie() = Trie{Any}()
-Trie(ks::AbstractVector{K}, vs::AbstractVector{V}) where {K<:AbstractString,V} = Trie{V}(ks, vs)
-Trie(kv::AbstractVector{Tuple{K,V}}) where {K<:AbstractString,V} = Trie{V}(kv)
-Trie(kv::AbstractDict{K,V}) where {K<:AbstractString,V} = Trie{V}(kv)
-Trie(ks::AbstractVector{K}) where {K<:AbstractString} = Trie{Nothing}(ks, similar(ks, Nothing))
+Trie() = Trie{Any,Any}()
+Trie(ks::AbstractVector{K}, vs::AbstractVector{V}) where {K,V} = Trie{eltype(K),V}(ks, vs)
+Trie(kv::AbstractVector{Tuple{K,V}}) where {K,V} = Trie{eltype(K),V}(kv)
+Trie(kv::AbstractDict{K,V}) where {K,V} = Trie{eltype(K),V}(kv)
+Trie(ks::AbstractVector{K}) where {K} = Trie{eltype(K),Nothing}(ks, similar(ks, Nothing))
 
-function Base.setindex!(t::Trie{T}, val, key::AbstractString) where T
-    value = convert(T, val) # we don't want to iterate before finding out it fails
+function Base.setindex!(t::Trie{K,V}, val, key) where {K,V}
+    value = convert(V, val) # we don't want to iterate before finding out it fails
     node = t
     for char in key
         if !haskey(node.children, char)
-            node.children[char] = Trie{T}()
+            node.children[char] = Trie{K,V}()
         end
         node = node.children[char]
     end
     node.is_key = true
     node.value = value
 end
 
-function Base.getindex(t::Trie, key::AbstractString)
+function Base.getindex(t::Trie, key)
     node = subtrie(t, key)
     if node != nothing && node.is_key
         return node.value
     end
     throw(KeyError("key not found: $key"))
 end
 
-function subtrie(t::Trie, prefix::AbstractString)
+function subtrie(t::Trie, prefix)
     node = t
     for char in prefix
         if !haskey(node.children, char)
@@ -66,30 +62,38 @@ function subtrie(t::Trie, prefix::AbstractString)
     return node
 end
 
-function Base.haskey(t::Trie, key::AbstractString)
+function Base.haskey(t::Trie, key)
     node = subtrie(t, key)
     node != nothing && node.is_key
 end
 
-function Base.get(t::Trie, key::AbstractString, notfound)
+function Base.get(t::Trie, key, notfound)
     node = subtrie(t, key)
     if node != nothing && node.is_key
         return node.value
     end
     return notfound
 end
 
-function Base.keys(t::Trie, prefix::AbstractString="", found=AbstractString[])
+_concat(prefix::String, char::Char) = string(prefix, char)
+_concat(prefix::Vector{T}, char::T) where {T} = vcat(prefix, char)
+
+_empty_prefix(::Trie{Char,V}) where {V} = ""
+_empty_prefix(::Trie{K,V}) where {K,V} = K[]
+
+function Base.keys(t::Trie{K,V},
+                   prefix=_empty_prefix(t),
+                   found=Vector{typeof(prefix)}()) where {K,V}
     if t.is_key
         push!(found, prefix)
     end
     for (char,child) in t.children
-        keys(child, string(prefix,char), found)
+        keys(child, _concat(prefix, char), found)
     end
     return found
 end
 
-function keys_with_prefix(t::Trie, prefix::AbstractString)
+function keys_with_prefix(t::Trie, prefix)
     st = subtrie(t, prefix)
     st != nothing ? keys(st,prefix) : []
 end
@@ -101,7 +105,7 @@ end
 # see the comments and implementation below for details.
 struct TrieIterator
     t::Trie
-    str::AbstractString
+    str
 end
 
 # At the start, there is no previous iteration,
@@ -120,11 +124,11 @@ function Base.iterate(it::TrieIterator, (t, i) = (it.t, 0))
     end
 end
 
-partial_path(t::Trie, str::AbstractString) = TrieIterator(t, str)
+partial_path(t::Trie, str) = TrieIterator(t, str)
 Base.IteratorSize(::Type{TrieIterator}) = Base.SizeUnknown()
 
 """
-    find_prefixes(t::Trie, str::AbstractString)
+    find_prefixes(t::Trie, str)
 
 Find all keys from the `Trie` that are prefix of the given string
 
@@ -137,10 +141,24 @@ julia> find_prefixes(t, "ABCDE")
  "A"
  "ABC"
  "ABCD"
+
+julia> t′ = Trie([1:1, 1:3, 1:4, 2:4]);
+
+julia> find_prefixes(t′, 1:5)
+3-element Vector{UnitRange{Int64}}:
+ 1:1
+ 1:3
+ 1:4
+
+julia> find_prefixes(t′, [1,2,3,4,5])
+3-element Vector{Vector{Int64}}:
+ [1]
+ [1, 2, 3]
+ [1, 2, 3, 4]
 ```
 """
-function find_prefixes(t::Trie, str::AbstractString)
-    prefixes = AbstractString[]
+function find_prefixes(t::Trie, str::T) where {T}
+    prefixes = T[]
     it = partial_path(t, str)
     idx = 0
     for t in it
@@ -150,4 +168,4 @@ function find_prefixes(t::Trie, str::AbstractString)
         idx = nextind(str, idx)
     end
     return prefixes
-end
+end

test/test_deprecations.jl

@@ -1,7 +1,7 @@
 # These are the tests for deprecated features, they should be deleted along with them
 
 @testset "Trie: path iterator" begin
-    t = Trie{Int}()
+    t = Trie{Char,Int}()
     t["rob"] = 27
     t["roger"] = 52
     t["kevin"] = Int8(11)

test/test_trie.jl

@@ -1,6 +1,6 @@
 @testset "Trie" begin
     @testset "Core Functionality" begin
-        t = Trie{Int}()
+        t = Trie{Char,Int}()
         t["amy"] = 56
         t["ann"] = 15
         t["emma"] = 30
@@ -19,14 +19,14 @@
         ks = ["amy", "ann", "emma", "rob", "roger"]
         vs = [56, 15, 30, 27, 52]
         kvs = collect(zip(ks, vs))
-        @test isa(Trie(ks, vs), Trie{Int})
-        @test isa(Trie(kvs), Trie{Int})
-        @test isa(Trie(Dict(kvs)), Trie{Int})
-        @test isa(Trie(ks), Trie{Nothing})
+        @test isa(Trie(ks, vs), Trie{Char,Int})
+        @test isa(Trie(kvs), Trie{Char,Int})
+        @test isa(Trie(Dict(kvs)), Trie{Char,Int})
+        @test isa(Trie(ks), Trie{Char,Nothing})
     end
 
     @testset "partial_path iterator" begin
-        t = Trie{Int}()
+        t = Trie{Char,Int}()
         t["rob"] = 27
         t["roger"] = 52
         t["kevin"] = Int8(11)
@@ -53,7 +53,7 @@
         @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")
@@ -66,4 +66,16 @@
         @test prefixes == ["東京都", "東京都渋谷区"]
     end
 
+    @testset "non-string indexing" begin
+        t = Trie{Int,Int}()
+        t[[1,2,3,4]] = 1
+        t[[1,2]] = 2
+        @test haskey(t, [1,2])
+        @test get(t, [1,2], nothing) == 2
+        st = subtrie(t, [1,2,3])
+        @test keys(st) == [[4]]
+        @test st[[4]] == 1
+        @test find_prefixes(t, [1,2,3,5]) == [[1,2]]
+        @test find_prefixes(t, 1:3) == [1:2]
+    end
 end # @testset Trie