Parent-child-sums
unknown
python
2 years ago
1.6 kB
12
Indexable
""" 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)
Editor is loading...