add quick2

6 files changed, 112 insertions, 8 deletions

diff --git a/Makefile b/Makefile
index 24dd61f..cf5cf2c 100644
--- a/Makefile
+++ b/Makefile
@@ -5,7 +5,7 @@ run:
 test:
 	go test ./... -v
 bench:
-	go test -run=10 -bench=. ./... -v
+	go test -run=xxx -bench=. ./... -v
 profile:
-	go test -bench=QuickSort ./sort -memprofile memprofile.out -cpuprofile profile.out
-	go tool pprof profile.out
+	go test -run=xxx -bench=Quick2Sort ./sort -memprofile memprofile.out -cpuprofile cpuprofile.out
+	go tool pprof cpuprofile.out
diff --git a/ds/elem.go b/ds/elem.go
index 0589e81..22bc9d9 100644
--- a/ds/elem.go
+++ b/ds/elem.go
@@ -1,10 +1,12 @@
 package ds
 
 type Elem interface {
-	Equal(a Elem) bool
-	Lower(a Elem) bool
-	LowerEqual(a Elem) bool
-	Higher(a Elem) bool
-	HigherEqual(a Elem) bool
+	Equal(e Elem) bool
+	Lower(e Elem) bool
+	LowerEqual(e Elem) bool
+	Higher(e Elem) bool
+	HigherEqual(e Elem) bool
+	Compare(e Elem) int
+	CompareCB(e Elem, lowerCB, higherCB, equalsCB func())
 	Int() int
 }
diff --git a/ds/integer.go b/ds/integer.go
index 4abb349..296e5bb 100644
--- a/ds/integer.go
+++ b/ds/integer.go
@@ -62,3 +62,26 @@ func (i Integer) Higher(j Elem) bool {
 func (i Integer) HigherEqual(j Elem) bool {
 	return i.Val >= j.Int()
 }
+
+func (i Integer) Compare(j Elem) int {
+    jVal := j.Int()
+    switch {
+      case i.Val < jVal:
+      return -1
+      case i.Val > jVal:
+      return 1
+    }
+	return 0
+}
+
+func (i Integer) CompareCB(j Elem, lowerCB, higherCB, equalsCB func()) {
+    jVal := j.Int()
+    switch {
+      case i.Val < jVal:
+        lowerCB()
+      case i.Val > jVal:
+        higherCB()
+      default:
+        equalsCB()
+    }
+}
diff --git a/ds/integer_test.go b/ds/integer_test.go
new file mode 100644
index 0000000..9be55eb
--- /dev/null
+++ b/ds/integer_test.go
@@ -0,0 +1,24 @@
+package ds
+
+import "testing"
+
+func TestCompare(t *testing.T) {
+	i := Integer{1}
+	j := Integer{10}
+	res := i.Compare(j)
+	if res != -1 {
+		t.Errorf("%v must be lower than %v, but got %v", i, j, res)
+	}
+
+	res = j.Compare(i)
+	if res != 1 {
+		t.Errorf("%v must be higher than %v, but got %v", j, i, res)
+	}
+
+	i = Integer{2}
+	j = Integer{2}
+	res = i.Compare(j)
+	if res != 0 {
+		t.Errorf("%v must be equal to %v, but got %v", i, j, res)
+	}
+}
diff --git a/sort/quick2.go b/sort/quick2.go
new file mode 100644
index 0000000..688e410
--- /dev/null
+++ b/sort/quick2.go
@@ -0,0 +1,43 @@
+package sort
+
+import (
+	"algorithms/ds"
+)
+
+// Quick2 uses a 3-way partitioning so it is more efficient
+// dealing with duplicates
+func Quick2(a ds.ArrayList) ds.ArrayList {
+	a = Shuffle(a)
+	quick2(a, 0, len(a)-1)
+	return a
+}
+
+func quick2(a ds.ArrayList, lo, hi int) {
+	if hi <= lo {
+		return
+	}
+
+	lt := lo    // Lower than
+	i := lo + 1 // lt..i contain duplicates
+	gt := hi    // Greater than
+	v := a[lo]  // Partitioning item
+
+	for i <= gt {
+		switch a[i].Compare(v) {
+		case -1:
+			a.Swap(lt, i)
+			lt++
+			i++
+		case 1:
+			a.Swap(i, gt)
+			gt--
+		default:
+			// Duplicate
+			i++
+		}
+	}
+	// Now a[lo..lt-1] < v = a[lt..gt] < a[gt+1..hi]
+
+	quick2(a, lo, lt-1)
+	quick2(a, gt+1, hi)
+}
diff --git a/sort/sort_test.go b/sort/sort_test.go
index a0e4729..3a44827 100644
--- a/sort/sort_test.go
+++ b/sort/sort_test.go
@@ -62,6 +62,12 @@ func TestQuickSort(t *testing.T) {
 	}
 }
 
+func TestQuick2Sort(t *testing.T) {
+	for i := 1; i <= maxLength; i *= 10 {
+		test(Quick2, i, t)
+	}
+}
+
 func TestShuffleSort(t *testing.T) {
 	for i := 10; i <= maxLength; i *= 10 {
 		testShuffle(Shuffle, i, t)
@@ -110,6 +116,12 @@ func BenchmarkQuickSort(b *testing.B) {
 	}
 }
 
+func BenchmarkQuick2Sort(b *testing.B) {
+	for i := 1; i <= maxLength; i *= 10 {
+		benchmark(Quick2, i, b)
+	}
+}
+
 func BenchmarkShuffleSort(b *testing.B) {
 	for i := 1; i <= maxLength; i *= 10 {
 		benchmark(Shuffle, i, b)