perf: create less BigDecimal in load balancing

TimefoldAI · Aug 13, 2024 · 9b67531 · 9b67531
1 parent 132a232
commit 9b67531
Showing 1 changed file with 10 additions and 7 deletions.
diff --git a/core/src/main/java/ai/timefold/solver/core/impl/score/stream/collector/LoadBalanceImpl.java b/core/src/main/java/ai/timefold/solver/core/impl/score/stream/collector/LoadBalanceImpl.java
@@ -96,13 +96,16 @@ public Map<Balanced_, Long> loads() {
     @Override
     public BigDecimal unfairness() {
         var totalToBalanceCount = balancedItemCountMap.size();
-        if (totalToBalanceCount == 0) {
-            return BigDecimal.ZERO;
-        }
-        return BigDecimal.valueOf(squaredDeviationFractionNumerator)
-                .divide(BigDecimal.valueOf(totalToBalanceCount), MathContext.DECIMAL32) // Compute w/ greater precision.
-                .add(BigDecimal.valueOf(squaredDeviationIntegralPart))
-                .sqrt(RESULT_MATH_CONTEXT);
+        return switch (totalToBalanceCount) {
+            case 0 -> BigDecimal.ZERO;
+            case 1 -> BigDecimal.valueOf(squaredDeviationFractionNumerator + squaredDeviationIntegralPart)
+                    .sqrt(RESULT_MATH_CONTEXT);
+            default -> { // Only do the final sqrt as BigDecimal, fast floating point math is good enough for the rest.
+                var tmp = (squaredDeviationFractionNumerator / (double) totalToBalanceCount) + squaredDeviationIntegralPart;
+                yield BigDecimal.valueOf(tmp)
+                        .sqrt(RESULT_MATH_CONTEXT);
+            }
+        };
     }
 
 }