diff --git a/vbt/vbt_test.go b/vbt/vbt_test.go index 587f86b..2e2677b 100644 --- a/vbt/vbt_test.go +++ b/vbt/vbt_test.go @@ -612,7 +612,7 @@ func BenchmarkPut(b *testing.B) { b.ResetTimer() b.StartTimer() - execCount := 1000 + execCount := 50 b.N = len(l) * execCount for i := 0; i < execCount; i++ { tree := New(compare.Int) diff --git a/vbtkey/vbtkey.go b/vbtkey/vbtkey.go index fd0c549..728f397 100644 --- a/vbtkey/vbtkey.go +++ b/vbtkey/vbtkey.go @@ -422,38 +422,49 @@ func (tree *Tree) Put(key, value interface{}) { return } - cur := tree.root - parent := cur.parent - child := -1 - - for { - - if cur == nil { - parent.children[child] = node - node.parent = parent - - fixed := parent.parent - fsize := getSize(fixed) - if fsize == 3 { - lefts, rigths := getChildrenSize(fixed) - tree.fix3Size(fixed, lefts, rigths) - } - return - } + for cur := tree.root; ; { if cur.size > 8 { - ls, rs := cur.children[0].size, cur.children[1].size factor := cur.size / 10 // or factor = 1 - if rs >= ls*2+factor || ls >= rs*2+factor { - tree.fixSize(cur, ls, rs) + if cur.children[1].size >= cur.children[0].size*2+factor || cur.children[0].size >= cur.children[1].size*2+factor { + tree.fixSize(cur) } } cur.size++ - parent = cur + c := tree.Compare(key, cur.key) - child = (c + 2) / 2 - cur = cur.children[child] + if c < 0 { + if cur.children[0] == nil { + cur.children[0] = node + node.parent = cur + + if cur.parent != nil && cur.parent.size == 3 { + if cur.parent.children[0] == nil { + tree.lrrotate3(cur.parent) + } else { + tree.rrotate3(cur.parent) + } + } + return + } + cur = cur.children[0] + } else { + if cur.children[1] == nil { + cur.children[1] = node + node.parent = cur + + if cur.parent != nil && cur.parent.size == 3 { + if cur.parent.children[1] == nil { + tree.rlrotate3(cur.parent) + } else { + tree.lrotate3(cur.parent) + } + } + return + } + cur = cur.children[1] + } } } @@ -756,6 +767,7 @@ func (tree *Tree) rrotate(cur *Node) { // mov.children[l]不可能为nil mov.children[l].parent = cur + cur.children[l] = mov.children[l] // 解决mov节点孩子转移的问题 @@ -809,6 +821,7 @@ func (tree *Tree) lrotate(cur *Node) { // mov.children[l]不可能为nil mov.children[l].parent = cur + cur.children[l] = mov.children[l] // 解决mov节点孩子转移的问题 @@ -851,10 +864,9 @@ func (tree *Tree) fixSizeWithRemove(cur *Node) { for cur != nil { cur.size-- if cur.size > 8 { - ls, rs := getChildrenSize(cur) factor := cur.size / 10 // or factor = 1 - if rs >= ls*2+factor || ls >= rs*2+factor { - tree.fixSize(cur, ls, rs) + if cur.children[1].size >= cur.children[0].size*2+factor || cur.children[0].size >= cur.children[1].size*2+factor { + tree.fixSize(cur) } } cur = cur.parent @@ -881,7 +893,7 @@ func (tree *Tree) fix3Size(cur *Node, lefts, rigths int) { } } -func (tree *Tree) fixSize(cur *Node, lefts, rigths int) { +func (tree *Tree) fix6Size(cur *Node, lefts, rigths int) { if lefts > rigths { l := cur.children[0] llsize, lrsize := getChildrenSize(l) @@ -901,6 +913,24 @@ func (tree *Tree) fixSize(cur *Node, lefts, rigths int) { } } +func (tree *Tree) fixSize(cur *Node) { + if cur.children[0].size > cur.children[1].size { + llsize, lrsize := getChildrenSize(cur.children[0]) + if lrsize > llsize { + tree.rlrotate(cur) + } else { + tree.rrotate(cur) + } + } else { + rlsize, rrsize := getChildrenSize(cur.children[1]) + if rlsize > rrsize { + tree.lrrotate(cur) + } else { + tree.lrotate(cur) + } + } +} + func output(node *Node, prefix string, isTail bool, str *string) { if node.children[1] != nil { @@ -956,10 +986,10 @@ func outputfordebug(node *Node, prefix string, isTail bool, str *string) { if node.parent == nil { parentv = "nil" } else { - parentv = spew.Sprint(node.key) + ":" + spew.Sprint(node.parent.value) + parentv = spew.Sprint(node.parent.key) } suffix += parentv + "|" + spew.Sprint(node.size) + ")" - *str += spew.Sprint(node.key) + ":" + spew.Sprint(node.value) + suffix + "\n" + *str += spew.Sprint(node.key) + suffix + "\n" if node.children[0] != nil { newPrefix := prefix diff --git a/vbtkey/vbtkey_test.go b/vbtkey/vbtkey_test.go index c771b52..62cc94b 100644 --- a/vbtkey/vbtkey_test.go +++ b/vbtkey/vbtkey_test.go @@ -625,16 +625,18 @@ func TestPutStable(t *testing.T) { } func BenchmarkGodsRBPut(b *testing.B) { - tree := redblacktree.NewWithIntComparator() - l := loadTestData() b.ResetTimer() b.StartTimer() - b.N = len(l) - for _, v := range l { - tree.Put(v, v) + execCount := 50 + b.N = len(l) * execCount + for i := 0; i < execCount; i++ { + tree := redblacktree.NewWithIntComparator() + for _, v := range l { + tree.Put(v, v) + } } }