Bank account data structure

Question

I need to model simple account data structure. The user should be able:

• to add a transaction
• get information about transaction
• take all transaction of some type
• sum a transaction and all its children

---

class AccountModel(object):
    def __init__(self):
        self.types = {}
        self.amounts = {}
        self.parent_to_chilren = {}
        self.children_to_parent = {}
        self.of_type = {}

    def add(self, id, amount, type, parentId=None):
        assert(
            not(id in self.amounts),
            'there is already a transaction with such id'
        )
        self.amounts[id] = amount
        self.of_type[id] = type

        lstOfType = self.types.get(type, [])
        lstOfType.append(id)
        self.types[type] = lstOfType

        if parentId:
            self.children_to_parent[id] = parentId
            lstOfChilren = self.parent_to_chilren.get(parentId, [])
            lstOfChilren.append(id)
            self.parent_to_chilren[parentId] = lstOfChilren

    def get(self, id):
        return (
            self.amounts[id], self.of_type[id],
            self.children_to_parent.get(id, None)
        )

    def get_list_by_type(self, type):
        return self.types[type]

    def sum(self, id):
        lst = self.parent_to_chilren.get(id, [])
        res = self.amounts[id]
        for i in lst:
            res += self.sum(i)
        return res

And the tests:

import unittest

from challenge import AccountModel
class AccountModelTestCase(unittest.TestCase):

    def setUp(self):
        self.model = AccountModel()

    def test_can_add_a_transaction_and_get_information(self):
        self.model.add(1, 10, 'car')
        self.assertEqual(
            self.model.get(1),
            (10, 'car', None),
        )

    def test_can_add_a_transaction_and_get_information_with_parent_id(self):
        self.model.add(1, 10, 'car')
        self.model.add(2, 10, 'car', 1)
        self.assertEqual(
            self.model.get(2),
            (10, 'car', 1),
        )

    def test_can_add_few_transactions_of_a_different_type_and_them(self):
        self.model.add(1, 10, 'car')
        self.model.add(2, 15, 'shopping')
        self.model.add(3, 18, 'car')
        self.assertItemsEqual(
            self.model.get_list_by_type('car'),
            [1, 3]
        )

    def test_can_add_few_transactions_of_a_different_type_and_calc_sum(self):
        self.model.add(1, 10, 'car')
        self.model.add(2, 15, 'shopping', 1)
        self.model.add(3, 18, 'car')
        self.assertEqual(
            self.model.sum(1),
            25
        )
        self.assertEqual(
            self.model.sum(2),
            15
        )
        self.assertEqual(
            self.model.sum(3),
            18
        )

    def test_can_add_few_transactions_parent_of_parent_of_parent(self):
        self.model.add(1, 10, 'car')
        self.model.add(2, 15, 'shopping', 1)
        self.model.add(3, 18, 'car', 2)
        self.assertEqual(
            self.model.sum(1),
            43
        )
        self.assertEqual(
            self.model.sum(2),
            33
        )


if __name__ == '__main__':
    unittest.main()

The questions:

• Is there a more effective (in terms of used memory) implementation?
• Is there a way to change a data structure and provide more effective (in terms of complexity) sum function implementation?

Answer 1

Your problem statement is a bit unclear,I'm uncertain what they mean by the children of a transition. In addition I would ask the interviewer what types (categories) of transaction it should support. Let's assume the basic ones, withdrawal, deposit and transfers.

I wouldn't store this information directly into the AccountModel, instead I would create a Transaction class which holds information regarding the transaction which is general and present for each type of transaction, namely: an id, amount, timestamp.

Then you can create subclasses Withdrawal, Deposit and Transfer, these can then be stored in a list of account which represents the history of the account's actions. An action is then performed by creating the action: account.perform(Withdrawal(10)), where accountis the an object of the AccountModelclass. In the perform()method I would check the type of transaction and update the account balance accordingly (which you store in the account directly). This way you don't need to use the sum()to run over the list of transactions each time.

Answer 2

Here are a few thoughts:

• Use snake_case for all variable names. This is the naming convention endorsed by PEP 8, the Python style guide.

  (If you can’t remember what the style is for an interview, at least be consistent – I see both snake_case and camelCase in this short script.

• Don’t use asserts as a substitute for exceptions. Asserts are for ensuring a program hasn’t reached an impossible state, whereas it is quite possible that a user could try to create two transactions with the same ID. You should raise an Exception instead.

  Supplementary question:

• Should the user be managing their own ids? If you want to enforce uniqueness, perhaps the AccountModel class should be managing ID numbers? Perhaps that’s a detail that a user shouldn’t worry about? How does a user know what’s a valid ID? etc.

  For the caller’s internal bookkeeping, you could create an ID for them, and return it from the add() method.

• Your data model is quite confusing. I made a quick diagram to follow where inputs to the add() method go. Here it is:

  

  There’s lots of duplicate information (which is quite memory inefficient), and it’s not immediately obvious which attributes I should be inspecting if I want information about a transaction.

  I’d suggest creating a flat dict of transactions, in which the keys are IDs (auto-generated by the class), and the values are namedtuples representing a single transaction. This is much more memory efficient.

  You then implement things like the transactions_of_type attribute by doing a list comprehension over this list.

• Use better variable names. Things like i, lst and res in the sum() method aren’t very meaningful.