This commit is contained in:
Nikolai Hartmann 2023-09-19 11:25:40 +02:00
parent 18479c5b2e
commit 798d068ec0

View file

@ -1,6 +1,6 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
from itertools import chain, combinations from itertools import chain, combinations
from typing import Generator from typing import Generator, Callable, Any
def zerosum_subgroups( def zerosum_subgroups(
@ -14,25 +14,34 @@ def zerosum_subgroups(
yield combination yield combination
def split_dict(d: dict, condition: Callable[[Any, Any], bool]) -> tuple[dict, dict]:
first = {}
second = {}
for k, v in d.items():
if condition(k, v):
first[k] = v
else:
second[k] = v
return first, second
def solve_greedily( def solve_greedily(
balances: dict[str, int], tolerance: int = 0 balances: dict[str, int], tolerance: int = 0
) -> dict[tuple[str, str], int]: ) -> dict[tuple[str, str], int]:
creditors = {} creditors, debitors = split_dict(balances, lambda k, v: v > 0)
debitors = {}
for k, v in balances.items():
if v > 0:
creditors[k] = v
else:
debitors[k] = v
transactions = {} transactions = {}
while not all( while not all(
abs(value) <= tolerance abs(value) <= tolerance
for value in chain(creditors.values(), debitors.values()) for value in chain(creditors.values(), debitors.values())
): ):
for debitor, debit_value in sorted(debitors.items(), key=lambda x: x[1]): for debitor, debit_value in sorted(
debitors.items(),
key=lambda x: x[1],
):
for creditor, credit_value in sorted( for creditor, credit_value in sorted(
creditors.items(), key=lambda x: x[1], reverse=True creditors.items(),
key=lambda x: x[1],
reverse=True,
): ):
sum_value = credit_value + debit_value sum_value = credit_value + debit_value
if abs(debit_value) <= credit_value: if abs(debit_value) <= credit_value:
@ -44,17 +53,15 @@ def solve_greedily(
debitors[debitor] = sum_value debitors[debitor] = sum_value
transactions[debitor, creditor] = credit_value transactions[debitor, creditor] = credit_value
break break
return transactions return transactions
def solve(balances: dict[str, int], tolerance: int = 0) -> dict[tuple[str, str], int]: def solve(balances: dict[str, int], tolerance: int = 0) -> dict[tuple[str, str], int]:
possibilities = [] possibilities = []
for subgroup in zerosum_subgroups(balances, tolerance): for subgroup in zerosum_subgroups(balances, tolerance):
transactions_sub = solve({k: balances[k] for k in subgroup}, tolerance) balances_sub, balances_other = split_dict(balances, lambda k, v: k in subgroup)
transactions_other = solve( transactions_sub = solve(balances_sub, tolerance)
{k: balances[k] for k in balances if not k in subgroup}, tolerance transactions_other = solve(balances_other, tolerance)
)
possibilities.append(transactions_sub | transactions_other) possibilities.append(transactions_sub | transactions_other)
if not possibilities: if not possibilities:
possibilities.append(solve_greedily(balances, tolerance)) possibilities.append(solve_greedily(balances, tolerance))