diff --git a/dynamic_programming/bitmask.py b/dynamic_programming/bitmask.py
index 56bb8e96ba02..c92d718858bd 100644
--- a/dynamic_programming/bitmask.py
+++ b/dynamic_programming/bitmask.py
@@ -12,7 +12,7 @@
 
 
 class AssignmentUsingBitmask:
-    def __init__(self, task_performed, total):
+    def __init__(self, task_performed: int, total: int):
         self.total_tasks = total  # total no of tasks (N)
 
         # DP table will have a dimension of (2^M)*N
@@ -27,7 +27,7 @@ def __init__(self, task_performed, total):
         # to 1
         self.final_mask = (1 << len(task_performed)) - 1
 
-    def count_ways_until(self, mask, task_no):
+    def count_ways_until(self, mask: int, task_no: int) -> int:
         # if mask == self.finalmask all persons are distributed tasks, return 1
         if mask == self.final_mask:
             return 1
@@ -60,7 +60,7 @@ def count_ways_until(self, mask, task_no):
 
         return self.dp[mask][task_no]
 
-    def count_no_of_ways(self, task_performed):
+    def count_no_of_ways(self, task_performed: int) -> int:
         # Store the list of persons for each task
         for i in range(len(task_performed)):
             for j in task_performed[i]: