Embercode
✕
Dynamic ProgrammingLesson 20 of 32

40 min read

Bitmask DP in Python: assigning tasks to people

Why this matters

A manager has 4 people and 4 tasks. Each person can do each task, but the cost differs. Person 0 is fast at task 1 and slow at task 3. Person 2 is the opposite. The manager wants one task per person, every task covered, and the smallest total cost. Some pairs may be impossible, for example when a person does not have the right licence for a task. This is the assignment problem with forbidden pairs.

With 4 people there are only 24 ways to hand out the tasks. That number grows quickly. With 10 people there are 10!, which is 3,628,800. With 20 people there are 20!, which is 2,432,902,008,176,640,000. Checking every assignment is not practical at that size.

Bitmask dynamic programming (DP) gives the answer by a different route. It processes the people in order, person 0 first, and for each set of tasks already handed out it stores the cheapest total cost. A set of tasks is stored as one whole number, a bitmask, which the previous lesson in this module, on the travelling salesman, explains in full. This lesson uses bitmasks in a different way: the mask is the state of the whole problem, not a route. A state is a snapshot of what has been decided so far.

The lesson covers two things. The first is how to fill the table and recover the assignment. The second is how to count the states, and how to check which ones can actually be reached. A reachable state is a set of tasks that some valid assignment of the first few people can produce. Counting reachable states tells you whether the problem has any solution at all, and how much of the table does real work.

This lesson assumes you know lists, for loops, if statements, functions and the bit operators &, | and ^. A list is an ordered row of values, written [10, 20, 30], and lst[i] reads position i, counting from 0. If the bit operators are new, read the travelling salesman lesson first, because it explains them in full.

The idea: the state is the set of tasks already handed out

Number the people 0 to 3 and the tasks 0 to 3. Suppose the first two people have already been given tasks. The tasks used so far form a set, say tasks 1 and 3. That set is the state. Two facts follow from it.

First, the number of tasks in the set equals the number of people already placed. Two tasks are used, so people 0 and 1 are placed, and person 2 is next. The program does not store the person number at all. It calls mask.bit_count() to find it.

Second, the state tells you which tasks are still free. Tasks not in the set are free for the remaining people. So the state contains everything the future choices depend on. As in the travelling salesman lesson, the cheapest way to reach a state is the only one that needs to be kept.

The table is best[mask]: the smallest total cost of giving the tasks in mask to the first mask.bit_count() people. Here best[0] is 0, because no one is placed and nothing is used. The answer is best[all], where all is the set of every task. In Python, all is a built-in function name, so the program calls the value all_tasks.

To fill the table, take a set mask with a known cost. The next person is p = mask.bit_count(). For each task t not in mask that person p can do, the set mask with t added is a new state. Its cost is the old cost plus cost[p][t]. Keep the smaller value when a state can be reached in more than one way.

Process the masks in increasing numeric order. Adding a task always makes the number bigger, so every smaller set is finished before a bigger set reads its value.

Worked example: the first three layers by hand

The cost table for the example. A row is a person, a column is a task:

           task 0  task 1  task 2  task 3
person 0      9       2       7       8
person 1      6       4       3       7
person 2      5       8       1       8
person 3      7       6       9       4

Every entry is possible in this example.

Layer 1, one person placed, so each set has one task:

  • {0}: 9, person 0 does task 0
  • {1}: 2, person 0 does task 1
  • {2}: 7, person 0 does task 2
  • {3}: 8, person 0 does task 3

Layer 2, two people placed. Person 1 is next, so the program adds a task using row 1 of the table. Each set has two ways to be reached, and the program keeps the smaller one:

  • {0, 1}: from {0} add task 1, 9 + 4 = 13; from {1} add task 0, 2 + 6 = 8. Keep 8.
  • {0, 2}: from {0} add task 2, 9 + 3 = 12; from {2} add task 0, 7 + 6 = 13. Keep 12.
  • {0, 3}: from {0} add task 3, 9 + 7 = 16; from {3} add task 0, 8 + 6 = 14. Keep 14.
  • {1, 2}: from {1} add task 2, 2 + 3 = 5; from {2} add task 1, 7 + 4 = 11. Keep 5.
  • {1, 3}: from {1} add task 3, 2 + 7 = 9; from {3} add task 1, 8 + 4 = 12. Keep 9.
  • {2, 3}: from {2} add task 3, 7 + 7 = 14; from {3} add task 2, 8 + 3 = 11. Keep 11.

Layer 3, three people placed. Person 2 is next, so row 2 is used. Each set has three predecessors:

  • {0, 1, 2}: 9, from {0, 1} adding task 2 (8 + 1)
  • {0, 1, 3}: 14, from {1, 3} adding task 0 (9 + 5)
  • {0, 2, 3}: 15, from {0, 3} adding task 2 (14 + 1)
  • {1, 2, 3}: 10, from {1, 3} adding task 2 (9 + 1)

Layer 4, all four tasks. Person 3 is last, so row 3 is used. The predecessor that gives the smallest total is {0, 1, 2}, from which person 3 takes task 3 at a cost of 9 + 4 = 13. The other three predecessors give 17, 21 and 23. So the best total is 13.

Recover the assignment by walking back from the full set, the same way the travelling salesman lesson rebuilt its route. Remember the task that was added last at each step:

  • Full set: the last task added was task 3, so person 3 does task 3.
  • Set {0, 1, 2}: the last task added was task 2, so person 2 does task 2.
  • Set {0, 1}: the last task added was task 0, so person 1 does task 0.
  • Set {1}: the last task added was task 1, so person 0 does task 1.

The assignment is person 0 to task 1, person 1 to task 0, person 2 to task 2, and person 3 to task 3. The total is 2 + 6 + 1 + 4 = 13.

The program

Create a file named assign.py:

Python · editable

Run it with python3 assign.py. The output is:

best total cost: 13
task for person 0..3: [1, 0, 2, 3]
reachable task sets: 16 of 16
task sets visited: 16
reachable by number of people placed: [1, 4, 6, 4, 1]
blocked case cost: 20
blocked case reachable task sets: 15 of 16

The best total is 13 and the assignment matches the hand walk above. The None in the cost table marks a forbidden pair. The program checks for None and skips that pair. The Java version of this lesson used -1 for the same job. Python has None, which means "no value", so it is the clearer choice. A -1 would look like a real cost to a reader and to the + and < operators.

Plan is a dataclass. A dataclass is a class that mostly holds data. The @dataclass line writes the setup code for the four fields, so the class only lists them. The type hints such as int | None say what each field should hold. Python does not check them while the program runs. They are for readers and for tools that check code before it runs.

The line reachable by number of people placed is the state count broken down by layer. With 4 people the layers are the sets of size 0, 1, 2, 3 and 4. The number of sets of size k out of 4 tasks is the binomial coefficient C(4, k), which gives 1, 4, 6, 4, 1. The reachable count is the same in every layer for this table, so every set can be produced.

The states_visited value counts the sets the loop actually processes. Sets that are unreachable are skipped, so this number is the amount of real work the program does.

Counting states and checking reachability

The table has one entry for each subset of tasks, so there are 2^n states. For 4 tasks that is 16. For 20 tasks it is 2^20, which is 1,048,576. Each state tries at most n new tasks, so the work is at most about 2^n * n. For 20 people that is about 21 million steps.

Not every state is reachable. A set is reachable when some valid way exists for the first k people to take exactly the tasks in that set, where k is the size of the set. If person 0 cannot do task 1, then the set {1} can never be reached, because the only way to fill one task is for person 0 to do it.

The reachable count is useful for two reasons. First, if the full set is unreachable, no complete assignment exists. Second, the count tells you how much of the table matters. Here is the blocked example, where person 0 cannot do task 1, and persons 1, 2 and 3 cannot do task 2:

           task 0  task 1  task 2  task 3
person 0      9       X       7       8
person 1      6       4       X       7
person 2      5       8       X       8
person 3      7       6       X       4

In this table, X marks a forbidden pair, which the program stores as None. The set {1} is unreachable, because only person 0 can take a single task and person 0 cannot do task 1. The program reports 15 reachable sets out of 16, so {1} is the only unreachable set. The cheapest full assignment is person 0 on task 2 at cost 7, person 1 on task 1 at cost 4, person 2 on task 0 at cost 5, and person 3 on task 3 at cost 4. The total is 7 + 4 + 5 + 4 = 20, which the program prints.

Example 2: when no assignment exists

When task 1 is forbidden for everyone, every set that contains task 1 is unreachable, and the full set is one of them. Here is that case with 3 people and 3 tasks. Save this file in the same folder as assign.py, because it imports solve from there:

from assign import solve
 
cost = [
    [4, None, 6],
    [3, None, 5],
    [2, None, 7],
]
plan = solve(cost)
print("cost:", plan.cost)
print("reachable task sets:", plan.reachable_states(), "of", 2 ** len(cost))
print("task sets visited:", plan.states_visited)

Run it with python3 impossible.py. The output is:

cost: None
reachable task sets: 4 of 8
task sets visited: 4

The cost is None, which the program uses to mean "no full assignment exists". The four reachable sets are the empty set, {0}, {2} and {0, 2}. Every set that contains task 1 is unreachable.

Example 3: 20 people

The next test uses 20 people and 20 tasks. The forbidden pairs follow a pattern: person p cannot do task t when (p + t) is a multiple of 5. Save this file in the same folder as assign.py as big2.py:

import time
 
from assign import solve
 
n = 20
cost = [
    [None if (p + t) % 5 == 0 else (p * 7 + t * 13) % 50 + 1 for t in range(n)]
    for p in range(n)
]
start = time.perf_counter()
plan = solve(cost)
ms = (time.perf_counter() - start) * 1000
print("cost:", plan.cost)
print("reachable task sets:", plan.reachable_states(), "of", 2**n)
print("elapsed ms:", round(ms))

Run it with python3 big2.py. On the machine used to test this lesson, with Python 3.14.4, the output was:

cost: 120
reachable task sets: 1048544 of 1048576
elapsed ms: 1228

Only 32 of the 1,048,576 sets are unreachable here, because most people can do most tasks. The elapsed time changes from run to run and between machines.

The memory for this table is small. The best and last_task lists each have 2^20 slots, and each slot holds one 8-byte reference. Each list therefore takes 8,388,608 bytes, about 8 MiB, and the total is about 16 MiB. The totals stored in best are integer objects, which take more memory than the slot that points to them, so the real total is a little higher.

Comparing the work with trying every assignment

Trying every assignment means building each complete plan and adding up its cost. The table method never builds a whole plan. It visits each set once and tries at most n tasks from it. The two counts grow in very different ways. The numbers below are worked out from the formulas, not measured:

  • 10 people: 10! is 3,628,800 plans. The table does at most 2^10 * 10, which is 10,240 steps.
  • 15 people: 15! is 1,307,674,368,000 plans. The table does at most 2^15 * 15, which is 491,520 steps.
  • 20 people: 20! is 2,432,902,008,176,640,000 plans. The table does at most 2^20 * 20, which is 20,971,520 steps.

For 4 people the difference is small. There are 24 plans, and the table has 16 sets. At that size the plain loop over plans is simpler to write, and it is a fine way to check the table on a small case.

The table trades time for memory. Its size is 2^n slots per list, so at some size memory becomes the limit instead of time. At 25 people each list has 33,554,432 slots, which is about 256 MiB at 8 bytes per slot, and the program needs two of them, so about 512 MiB. The work at 25 people is 2^25 * 25, which is 838,860,800 steps. On a machine with little memory, the lists are the part that fails first.

Common mistakes

Forgetting the forbidden check. If the loop does not skip None costs, the program tries to add None to a number. Python stops with an error at once. Removing the check from the program above and running the blocked example prints this error:

TypeError: unsupported operand type(s) for +: 'int' and 'NoneType'

The Java version of this mistake did not stop. It printed an answer that was too small. Python's error is easier to find, but it still means the check is missing, so fix the check instead of catching the error.

Comparing None with a number. The result plan.cost is None when no assignment exists. A comparison such as plan.cost < 5 raises an error in that case:

TypeError: '<' not supported between instances of 'NoneType' and 'int'

Check for None before you use the value: if plan.cost is None:.

Making the table too small. The same size mistake from the travelling salesman lesson applies here. Both best and last_task need 2^n slots, so that the set number is a valid index.

Adding to an unreachable entry. The table starts every unreachable entry at float("inf"). Adding to infinity gives infinity, so the value stays unreachable, and Python integers never wrap around. The program still checks for infinity before it adds, because the check also avoids doing work that cannot change the answer.

Mixing up the set and the person. A state is the set of tasks, not the pair of person and task. The person is always mask.bit_count(). Storing the person separately adds a dimension to the table without adding information, and it can disagree with the set if one is updated and the other is not.

Practice

Exercise 1. A set of tasks is written as 0b0110. Which tasks are in the set, and how many people have been placed when the program reaches this state?

Show the answer

Bits 1 and 2 are set, so the tasks are 1 and 2. The number of people placed is bit_count of 6, which is 2.

Exercise 2. For n people, how many states does the table have? For n = 3, list the number of states at each layer, and check that they add up to the total.

Show the answer

There are 2^n states, so 8 for n = 3. The layers are the sets of size 0, 1, 2 and 3, with counts C(3, k): 1, 3, 3 and 1. They add up to 8.

Exercise 3. Using the 4 by 4 example table in this lesson, which person is placed last in the optimal assignment, and what does that person do?

Show the answer

Person 3 is placed last, because the program fills layer 4 from row 3 of the table. In the optimal assignment, person 3 does task 3 at cost 4.

Exercise 4. In a 5 by 5 table, person 0 can do only task 4. Every other person can do every task. How many sets of size 1 are reachable?

Show the answer

Exactly one. Only person 0 can take a single task, and the only task person 0 can do is task 4. So the only reachable set of size 1 is {4}.

Exercise 5. How many bytes does the best list take for 16 people? Use 8 bytes per slot, and ignore the integer objects and the list header.

Show the answer

The list has 2^16 = 65,536 slots. At 8 bytes each, that is 524,288 bytes, or 512 KiB.

Summary

  • The state is the set of tasks already handed out, stored as a bitmask. The person to place next is the number of tasks in the set.
  • The table entry for a set is the cheapest total cost for the people placed so far to take exactly those tasks.
  • There are 2^n states, but only the reachable ones hold a real value. Counting the reachable states per layer shows which sets can be produced.
  • If the full set is unreachable, no complete assignment exists. The program reports this as None.
  • The work is at most about 2^n * n steps, which made the 20-person test run in about a second on the test machine.
Check yourself: Bitmask assignment (6 questions) →Practice this topic: coding problems for Dynamic Programming →

Stuck? Ask the tutor

Ask about anything in this lesson. Answers are short and use simple words.

Sign in to ask the tutor.