wallet_groups.py

#!/usr/bin/env python3
from typing import List


class Wallet:
    def __init__(self, name: str):
        self.name = name

    def __repr__(self):
        return self.name


class WalletGroup:
    def __init__(self, wallets: List[Wallet]):
        self.wallets = wallets
        self.wallets_set = set([str(wallet) for wallet in wallets])

    def wallet_names(self) -> List[str]:
        return [wallet.name for wallet in self.wallets]


class WalletPair:
    def __init__(self, left_wallet: Wallet, right_wallet: Wallet):
        self.left = left_wallet
        self.right = right_wallet


def check_wallet_equality(
        pairs: List[WalletPair],
        left_wallet: WalletGroup,
        right_wallet: WalletGroup,
        ) -> bool:
    """
    Given a set of transactions (WalletPair objects) between 2 Wallets
    and 2 WalletGroups, we want to identify if these WalletGroups are the
    same cluster.

    We create an index where any Wallet maps to a set of all connected wallets.

    We create the set by first adding the WalletGroup values, then we add on the
    relationships from the WalletPair transactions.

    Finally, we query our datastructure to check equality

        Note: comparing sets is expensive, so if we the sets were very large we
        could optionally go over our final index and hash all set values to strings,
        making the equality comparison faster. However this would be a poor choice if
        the average set size is small and the number of keys is very high.
    """
    index = {}
    # First we populate the index with the WalletGroups
    for name in left_wallet.wallet_names():
        index[name] = left_wallet.wallets_set
    for name in right_wallet.wallet_names():
        index[name] = right_wallet.wallets_set

    # Then go over the pairs and add those relationships
    # to the index set
    for pair in pairs:
        left, right = pair.left, pair.right
        default_set = set([left.name, right.name])
        left_index = index.get(left.name, default_set)
        right_index = index.get(right.name, default_set)

        # Now merge them and update the values
        merged_index = right_index.union(left_index)
        index[left.name] = merged_index
        index[right.name] = merged_index

    # Now to check if we have equal wallet groups, just match
    # any name from the left and the right wallets in our index
    # and compare the sets
    key_left = left_wallet.wallet_names()[0]
    key_right = right_wallet.wallet_names()[0]
    return index[key_left] == index[key_right]


def run_tests():
    test_cases = [
        (
            [
                WalletPair(Wallet('A'), Wallet('B')),
                WalletPair(Wallet('C'), Wallet('D')),
            ],
            WalletGroup([Wallet('A'), Wallet('C')]),
            WalletGroup([Wallet('B'), Wallet('C')]),
            True,
        ),
        (
            [
                WalletPair(Wallet('A'), Wallet('B')),
                WalletPair(Wallet('C'), Wallet('D')),
            ],
            WalletGroup([Wallet('A'), Wallet('C')]),
            WalletGroup([Wallet('A'), Wallet('C')]),
            True,
        ),
        (
            [
                WalletPair(Wallet('A'), Wallet('B')),
                WalletPair(Wallet('C'), Wallet('D')),
            ],
            WalletGroup([Wallet('A'), Wallet('B')]),
            WalletGroup([Wallet('C'), Wallet('D')]),
            False,
        ),
    ]
    i = 0
    for case in test_cases:
        i += 1
        print(f"Running test case {i}")
        assert case[3] == check_wallet_equality(case[0], case[1], case[2])
    pass


if __name__ == '__main__':
    print("Solving wallets")
    run_tests()