diff --git a/src/trie.jl b/src/trie.jl index 03b1f933..2d827b8e 100644 --- a/src/trie.jl +++ b/src/trie.jl @@ -47,6 +47,27 @@ function Base.getindex(t::Trie, key) throw(KeyError("key not found: $key")) end +function Base.delete!(t::Trie, key; at_key=true) + if isempty(key) || (at_key && !haskey(t, key)) + return t + end + node = subtrie(t, key) + isnothing(node) && return t # no subtrie in t at key, do nothing + if at_key # only set is_key to false on the inner-most node + node.is_key = false + end + all_dangling = true # assume all children are dangling + for (k,v) in node.children # clean up dangling nodes + if isempty(v.children) && !v.is_key # dangling + delete!(node.children, k) + else + all_dangling = false # at least one not dangling + end + end + # only recurse if all children were dangling + return all_dangling ? delete!(t, @view(key[1:end-1]); at_key=false) : t +end + function subtrie(t::Trie, prefix) node = t for char in prefix diff --git a/test/test_trie.jl b/test/test_trie.jl index bce5ca52..6ba816b8 100644 --- a/test/test_trie.jl +++ b/test/test_trie.jl @@ -13,6 +13,22 @@ @test sort(keys(t)) == ["amy", "ann", "emma", "kevin", "rob", "roger"] @test t["rob"] == 27 @test sort(keys_with_prefix(t,"ro")) == ["rob", "roger"] + + delete!(t, "roger") + @test !haskey(t, "roger") + @test haskey(t, "rob") + @test get(t,"rob",nothing) == 27 + @test sort(keys(t)) == ["amy", "ann", "emma", "kevin", "rob"] + @test t["rob"] == 27 + @test sort(keys_with_prefix(t,"ro")) == ["rob"] + + # Ensure deletion cleans up dangling nodes + t = Trie(["A", "ABC"]) + delete!(t, "ABC") + t_no_abc = Trie(["A"]) + partial_no_abc = collect(partial_path(t_no_abc, "ABC")) + partial_original = collect(partial_path(t, "ABC")) + @test length(partial_no_abc) == length(partial_original) end @testset "Constructors" begin