goimagehash: Fix quickselect algorithm to pick median value.

corona10 · May 8, 2018 · 1903317 · 1903317
1 parent ad2b6de
commit 1903317
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 28 deletions.
diff --git a/etcs/utils.go b/etcs/utils.go
@@ -24,35 +24,38 @@ func MeanOfPixels(pixels []float64) float64 {
 func MedianOfPixels(pixels []float64) float64 {
 	tmp := make([]float64, len(pixels))
 	copy(tmp, pixels)
-	l := len(tmp) - 1
+	l := len(tmp)
 	pos := l / 2
-	v := quickSelect(tmp, 0, l, pos)
+	v := quickSelectMedian(tmp, 0, l-1, pos)
 	return v
 }
 
-func quickSelect(sequence []float64, low int, hi int, k int) float64 {
-	if hi-low <= 1 {
+func quickSelectMedian(sequence []float64, low int, hi int, k int) float64 {
+	if low == hi {
 		return sequence[k]
 	}
-	j := low
-	sequence[j], sequence[k] = sequence[k], sequence[j]
-	j++
-	for i := j; i < hi; i++ {
-		if sequence[i] < sequence[low] {
-			sequence[j], sequence[i] = sequence[i], sequence[j]
-			j++
-		}
-	}
-	j--
-	sequence[j], sequence[low] = sequence[low], sequence[j]
 
-	if k < j {
-		return quickSelect(sequence, low, j, k)
+	for low < hi {
+		pivot := low/2 + hi/2
+		pivotValue := sequence[pivot]
+		storeIdx := low
+		sequence[pivot], sequence[hi] = sequence[hi], sequence[pivot]
+		for i := low; i < hi; i++ {
+			if sequence[i] < pivotValue {
+				sequence[storeIdx], sequence[i] = sequence[i], sequence[storeIdx]
+				storeIdx++
+			}
+		}
+		sequence[hi], sequence[storeIdx] = sequence[storeIdx], sequence[hi]
+		if k <= storeIdx {
+			hi = storeIdx
+		} else {
+			low = storeIdx + 1
+		}
 	}
 
-	if k > j {
-		return quickSelect(sequence, j+1, hi, k-j)
+	if len(sequence)%2 == 0 {
+		return sequence[k-1]/2 + sequence[k]/2
 	}
-
-	return sequence[j]
+	return sequence[k]
 }
diff --git a/etcs/utils_test.go b/etcs/utils_test.go
@@ -30,9 +30,9 @@ func TestMedianPixels(t *testing.T) {
 	}{
 		{[]float64{0, 0, 0, 0}, 0},
 		{[]float64{1}, 1},
-		{[]float64{1, 2, 3, 4}, 2},
+		{[]float64{1, 2, 3, 4}, 2.5},
 		{[]float64{5, 3, 1, 7, 9}, 5},
-		{[]float64{98.3, 33.4, 105.44, 1500.4, 22.5, 66.6}, 98.3},
+		{[]float64{98.3, 33.4, 105.44, 1500.4, 22.5, 66.6}, 82.44999999999999},
 	} {
 		pixels := tt.pixels
 		result := MedianOfPixels(pixels)

diff --git a/hashcompute_test.go b/hashcompute_test.go
@@ -41,11 +41,11 @@ func TestHashCompute(t *testing.T) {
 		{"_examples/sample2.jpg", "_examples/sample2.jpg", PerceptionHash, "PerceptionHash", 0},
 		{"_examples/sample3.jpg", "_examples/sample3.jpg", PerceptionHash, "PerceptionHash", 0},
 		{"_examples/sample4.jpg", "_examples/sample4.jpg", PerceptionHash, "PerceptionHash", 0},
-		{"_examples/sample1.jpg", "_examples/sample2.jpg", PerceptionHash, "PerceptionHash", 34},
-		{"_examples/sample1.jpg", "_examples/sample3.jpg", PerceptionHash, "PerceptionHash", 7},
-		{"_examples/sample1.jpg", "_examples/sample4.jpg", PerceptionHash, "PerceptionHash", 31},
-		{"_examples/sample2.jpg", "_examples/sample3.jpg", PerceptionHash, "PerceptionHash", 31},
-		{"_examples/sample2.jpg", "_examples/sample4.jpg", PerceptionHash, "PerceptionHash", 23},
+		{"_examples/sample1.jpg", "_examples/sample2.jpg", PerceptionHash, "PerceptionHash", 32},
+		{"_examples/sample1.jpg", "_examples/sample3.jpg", PerceptionHash, "PerceptionHash", 2},
+		{"_examples/sample1.jpg", "_examples/sample4.jpg", PerceptionHash, "PerceptionHash", 30},
+		{"_examples/sample2.jpg", "_examples/sample3.jpg", PerceptionHash, "PerceptionHash", 34},
+		{"_examples/sample2.jpg", "_examples/sample4.jpg", PerceptionHash, "PerceptionHash", 20},
 	} {
 		file1, err := os.Open(tt.img1)
 		if err != nil {