Python org cost rollup problem

org_cost_rollup.py

# Org Cost Rollup
# Problem Statement
# You are given a company’s organization represented as a tree.
# You have n employees

# Each employee:
# has a unique integer ID 0 .. n-1
# has at most one manager
# may own cloud accounts that generate a direct monthly cost

# The CEO has no manager (id = -1).

# Input
# n – number of employees
# manager[i] – manager of employee i, or -1 if i has no manager
# direct_cost[i] – direct cloud cost of employee i

# queries – a list of employee IDs; for each employee x, return the total cost of the org subtree rooted at x.
# The total org cost for an employee x is:

# total_cost(x) = sum of direct_cost[i] for all employees i in x’s subtree (including x)

# Example Walkthrough
# n = 5
#                 0  1. 2. 3. 4
# manager     = [-1, 0, 0, 1, 1]

#                 0  1. 2. 3. 4
# direct_cost = [ 8, 3, 2, 5, 1 ]

# queries     = [0, 1, 3]

# org chart:
#         0(cost(8))
#       /   \
#     1     2.      
#   /   \
#   3     4
 
# employee.   subtree.      cost
# 0    {0, 1, 2, 3, 4}    8 + 3 + 2 + 5 + 1 = 19
# 1    {1, 3, 4}    3 + 5 + 1 = 9
# 3    {3}    5

# Output: [19, 9, 5]

def org_cost_rollup(
            self,
            n: int,
            manager: list[int],
            direct_cost: list[int],
            queries: list[int]
        ) -> list[int]: 
            
            # init
            total_cost = 0
            emps = []
            costs = []
            queue = []

            for mgr in queries :
                queue.append(mgr)

                while queue:
                    cur = queue.pop(0)
                    emps.append(cur)
                    
                    for i in range(len(manager)):
                        if manager[i] == cur:
                            queue.append(i)

                for emp in emps:
                    emp_cost = direct_cost[emp]
                    total_cost = total_cost + emp_cost

                costs.append(total_cost)
                total_cost = 0  # reset total_cost
                emps = []

            return costs