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?