diff --git a/dynamic_programming/bitmask.py b/dynamic_programming/bitmask.py index 4737a3419e8e..a4724448d1ed 100644 --- a/dynamic_programming/bitmask.py +++ b/dynamic_programming/bitmask.py @@ -62,6 +62,26 @@ def count_ways_until(self, mask, task_no): return self.dp[mask][task_no] def count_no_of_ways(self, task_performed): + """ + Calculate the total number of ways to assign tasks to people. + + Args: + task_performed: List where task_performed[i] contains task IDs + that person i can perform + + Returns: + Total number of valid task assignment ways + + >>> tp = [[1, 3, 4], [1, 2, 5], [3, 4]] + >>> AssignmentUsingBitmask(tp, 5).count_no_of_ways(tp) + 10 + >>> tp2 = [[1, 2], [1, 3]] + >>> AssignmentUsingBitmask(tp2, 3).count_no_of_ways(tp2) + 2 + >>> tp3 = [[1], [2]] + >>> AssignmentUsingBitmask(tp3, 2).count_no_of_ways(tp3) + 1 + """ # Store the list of persons for each task for i in range(len(task_performed)): for j in task_performed[i]: