Fix regression and more optimizations

flanglet · Jan 24, 2024 · 05439b6 · 05439b6
1 parent 1923111
commit 05439b6
Showing 1 changed file with 15 additions and 6 deletions.
diff --git a/v2/entropy/EntropyUtils.go b/v2/entropy/EntropyUtils.go
@@ -186,10 +186,6 @@ func NormalizeFrequencies(freqs []int, alphabet []int, totalFreq, scale int) (in
 			continue
 		}
 
-		if f > freqs[idxMax] {
-			idxMax = i
-		}
-
 		sf := int64(freqs[i]) * int64(scale)
 		var scaledFreq int
 
@@ -211,6 +207,10 @@ func NormalizeFrequencies(freqs []int, alphabet []int, totalFreq, scale int) (in
 		alphabetSize++
 		sumScaledFreq += scaledFreq
 		freqs[i] = scaledFreq
+
+		if scaledFreq > freqs[idxMax] {
+			idxMax = i
+		}
 	}
 
 	if alphabetSize == 0 {
@@ -224,22 +224,31 @@ func NormalizeFrequencies(freqs []int, alphabet []int, totalFreq, scale int) (in
 
 	if sumScaledFreq != scale {
 		delta := sumScaledFreq - scale
+errThr := freqs[idxMax] >> 4
 		var inc, absDelta int
 
-		if sumScaledFreq < scale {
+		if delta < 0 {
 			absDelta = -delta
 			inc = 1
 		} else {
 			absDelta = delta
 			inc = -1
 		}
 
-		if absDelta*10 < freqs[idxMax] {
+		if absDelta <= errThr {
 			// Fast path (small error): just adjust the max frequency
 			freqs[idxMax] -= delta
 			return alphabetSize, nil
 		}
 
+		if delta < 0 {
+			freqs[idxMax] += errThr 
+			sumScaledFreq += errThr
+		} else {
+			freqs[idxMax] -= errThr
+			sumScaledFreq -= errThr
+		}
+
 		// Slow path: spread error across frequencies
 		queue := make(sortByFreq, alphabetSize)
 		n := 0