cleanup and simplify subset function

This commit is contained in:
Nikolai Hartmann 2023-09-19 09:24:08 +02:00
parent cdca6df821
commit c9164a080a

View file

@ -2,25 +2,13 @@
from itertools import chain, combinations
from typing import Generator
balances = {
# "A": -20,
# "B": 10,
# "C": 3,
# "D": -43,
# should be possible with 3 transactions (A, B, C balance excactly)
"A": 50,
"B": -30,
"C": -20,
"D": -40,
}
balances["E"] = -sum(balances.values())
def find_zerosum_subgroups(
def zerosum_subgroups(
balances: dict[str, int],
) -> Generator[tuple[str, ...], None, None]:
for n in range(2, len(balances)):
for combination in combinations(balances, n):
if len(balances) < 3:
return
for combination in combinations(balances, len(balances) - 2):
if sum(balances[key] for key in combination) == 0:
yield combination
@ -34,7 +22,7 @@ def solve_greedily(balances: dict[str, int]) -> dict[tuple[str, str], int]:
else:
debitors[k] = v
txn = {}
transactions = {}
while not all(value == 0 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(
@ -44,25 +32,34 @@ def solve_greedily(balances: dict[str, int]) -> dict[tuple[str, str], int]:
if abs(debit_value) <= credit_value:
del debitors[debitor]
creditors[creditor] = sum_value
txn[debitor, creditor] = abs(debit_value)
transactions[debitor, creditor] = abs(debit_value)
else:
debitors[debitor] = sum_value
del creditors[creditor]
txn[debitor, creditor] = credit_value
debitors[debitor] = sum_value
transactions[debitor, creditor] = credit_value
break
return txn
return transactions
def solve(balances: dict[str, int]) -> dict[tuple[str, str], int]:
possibilities = []
for subgroup in find_zerosum_subgroups(balances):
txn_sub = solve({k: balances[k] for k in subgroup})
txn_other = solve({k: balances[k] for k in balances if not k in subgroup})
possibilities.append(txn_sub | txn_other)
for subgroup in zerosum_subgroups(balances):
transactions_sub = solve({k: balances[k] for k in subgroup})
transactions_other = solve({k: balances[k] for k in balances if not k in subgroup})
possibilities.append(transactions_sub | transactions_other)
if not possibilities:
possibilities.append(solve_greedily(balances))
return min(possibilities, key=lambda x: len(x))
print(solve(balances))
if __name__ == "__main__":
# should be possible with 3 transactions (A, B, C balance excactly)
balances = {
"A": 50,
"B": -30,
"C": -20,
"D": -40,
}
balances["E"] = -sum(balances.values())
print(solve(balances))