Parent-child-sums

"""
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)