Parent-child-sums

unknown
python
2 years ago
1.6 kB
8
Indexable
Never
```"""
Depth first search; parent child gathering

@author: ivbarrie
"""

from typing import List, Dict, Tuple
import unittest

def dfs(child_list: List[int], parent_child_pairs: List[Tuple[int, int]], parent: int) -> List[int]:
"""Depth first search for all descendants of parent."""
for p, c in parent_child_pairs:
if p == parent:
child_list.append(c)
dfs(child_list, parent_child_pairs, parent=c)

return child_list

def compute_total_vals(employees_to_mgrs: List[List[int]], employee_vals: Dict[int, int]) -> Dict[int, int]:
"""Compute sum of values including descendants."""
employee_total_vals_map = {}
employee_list = employees_to_mgrs[0]
mgrs_list = employees_to_mgrs[1]
id_set = set(employee_list) | set(mgrs_list)
mgr_employee_pairs = list(zip(mgrs_list, employee_list))  # could also use namedtuples
for _id in id_set:
total = employee_vals[_id] if _id in employee_vals else 0
descendants = dfs(child_list=[], parent_child_pairs=mgr_employee_pairs, parent=_id)
total += sum([employee_vals[d] for d in descendants])
employee_total_vals_map[_id] = total

return employee_total_vals_map

class Test(unittest.TestCase):

def test_basic_results(self):
employees_to_mgrs = [[1, 2, 3, 5, 6, 7], [4, 4, 2, 2, 4, 3]]
employee_vals = {idx: idx for idx in employees_to_mgrs[0]}
expected = {1: 1, 2: 17, 3: 10, 4: 24, 5: 5, 6: 6, 7: 7}
result = compute_total_vals(employees_to_mgrs=employees_to_mgrs, employee_vals=employee_vals)
self.assertGreater(result[4], result[1])
self.assertEqual(result, expected)

```