diff --git a/trie/trie.go b/trie/trie.go index 81766e8c..f78718fd 100644 --- a/trie/trie.go +++ b/trie/trie.go @@ -14,6 +14,7 @@ type Node struct { squash bool + // TODO inner node count. fix it NodeCnt int } @@ -26,7 +27,7 @@ const leafBranch = -1 // key. func NewTrie(keys [][]byte, values interface{}, squash bool) (root *Node, err error) { - root = &Node{Children: make(map[int]*Node), Step: 1, squash: squash} + root = &Node{Children: make(map[int]*Node), Step: 1, squash: squash, NodeCnt: 1} if keys == nil { return @@ -81,6 +82,60 @@ func (r *Node) Squash() int { return cnt } +// removeSameLeaf removes leaf that has the same value as preceding leaf. +// +// a ------->e =1 +// `>b------>f =2 +// `>c->d->g =2 // "g" and "d" is removed, c has other child and is kept. +// `--->h =3 +// +// Since 0.5.5 +func (r *Node) removeSameLeaf() { + + var prevValue interface{} = nil + + // wrapped as a generalized tree + s := &trieStringly{tnode: r} + + DepthFirst(s, + func(t Tree, parent, branch, node interface{}) { + + n := node.(*Node) + needRemove := false + + v, isLeaf := t.LeafVal(node) + if isLeaf { + if v == prevValue { + // same value no need to store + needRemove = true + } else { + prevValue = v + } + } else { + if len(n.Branches) == 0 { + needRemove = true + } + } + + if needRemove && parent != nil && branch != nil { + p := parent.(*Node) + b := branch.(int) + + delete(p.Children, b) + + for i, bb := range p.Branches { + if bb == b { + p.Branches = append(p.Branches[:i], p.Branches[i+1:]...) + } + } + if !isLeaf { + r.NodeCnt-- + } + + } + }) +} + // Search for `key` in a Trie. // // It returns 3 values of: diff --git a/trie/trie_test.go b/trie/trie_test.go index d104985e..4c8f0b56 100644 --- a/trie/trie_test.go +++ b/trie/trie_test.go @@ -6,6 +6,7 @@ import ( "github.com/openacid/errors" "github.com/openacid/slim/benchhelper" + "github.com/stretchr/testify/require" ) func TestTrie(t *testing.T) { @@ -527,3 +528,43 @@ func TestToStrings(t *testing.T) { t.Fatalf("expect: \n%v\n; but: \n%v\n", expect, trie.String()) } } + +func TestTrie_removeSameLeaf(t *testing.T) { + + ta := require.New(t) + + var keys = [][]byte{ + {'a', 'b', 'c'}, + {'a', 'b', 'c', 'd'}, + {'a', 'b', 'd'}, + {'a', 'b', 'd', 'e'}, + {'b', 'c'}, + {'b', 'c', 'd'}, + {'b', 'c', 'd', 'e'}, + {'c', 'd', 'e'}, + } + var values = []int{0, 0, 0, 3, 4, 5, 5, 5} + + want := ` +*2 +-097-> + -098->*2 + -099-> + -00$->=0 + -100-> + -101-> + -00$->=3 +-098-> + -099->*2 + -00$->=4 + -100-> + -00$->=5`[1:] + + trie, err := NewTrie(keys, values, false) + ta.Nil(err) + + trie.removeSameLeaf() + + ta.Equal(want, trie.String()) + ta.Equal(9, trie.NodeCnt, "non-leaf node count") +}