introduce tolerance

splitbill.py

@@ -5,15 +5,18 @@ from typing import Generator
 
 def zerosum_subgroups(
     balances: dict[str, int],
+    tolerance: int = 0,
 ) -> Generator[tuple[str, ...], None, None]:
     if len(balances) < 3:
         return
     for combination in combinations(balances, len(balances) - 2):
-        if sum(balances[key] for key in combination) == 0:
+        if abs(sum(balances[key] for key in combination)) <= tolerance:
             yield combination
 
 
-def solve_greedily(balances: dict[str, int]) -> dict[tuple[str, str], int]:
+def solve_greedily(
+    balances: dict[str, int], tolerance: int = 0
+) -> dict[tuple[str, str], int]:
     creditors = {}
     debitors = {}
     for k, v in balances.items():
@@ -23,7 +26,10 @@ def solve_greedily(balances: dict[str, int]) -> dict[tuple[str, str], int]:
             debitors[k] = v
 
     transactions = {}
-    while not all(value == 0 for value in chain(creditors.values(), debitors.values())):
+    while not all(
+        abs(value) <= tolerance
+        for value in chain(creditors.values(), debitors.values())
+    ):
         for debitor, debit_value in sorted(debitors.items(), key=lambda x: x[1]):
             for creditor, credit_value in sorted(
                 creditors.items(), key=lambda x: x[1], reverse=True
@@ -42,17 +48,29 @@ def solve_greedily(balances: dict[str, int]) -> dict[tuple[str, str], int]:
     return transactions
 
 
-def solve(balances: dict[str, int]) -> dict[tuple[str, str], int]:
+def solve(balances: dict[str, int], tolerance: int = 0) -> dict[tuple[str, str], int]:
     possibilities = []
-    for subgroup in zerosum_subgroups(balances):
+    for subgroup in zerosum_subgroups(balances, tolerance):
-        transactions_sub = solve({k: balances[k] for k in subgroup})
+        transactions_sub = solve({k: balances[k] for k in subgroup}, tolerance)
-        transactions_other = solve({k: balances[k] for k in balances if not k in subgroup})
+        transactions_other = solve(
+            {k: balances[k] for k in balances if not k in subgroup}, tolerance
+        )
         possibilities.append(transactions_sub | transactions_other)
     if not possibilities:
-        possibilities.append(solve_greedily(balances))
+        possibilities.append(solve_greedily(balances, tolerance))
     return min(possibilities, key=lambda x: len(x))
 
 
+def perform_transfers(
+    balances: dict[str, int], transactions: dict[tuple[str, str], int]
+) -> dict[str, int]:
+    balances = balances.copy()
+    for (sender, recipient), value in transactions.items():
+        balances[sender] += value
+        balances[recipient] -= value
+    return balances
+
+
 if __name__ == "__main__":
     # should be possible with 3 transactions (A, B, C balance excactly)
     balances = {