Title: Use a comprehension to make a two-dimensional list of lists in Python

[A two-dimensional list of lists]

You may have noticed that I've been writing a lot of graphics examples lately. I do that a lot when I'm feeling tired or stressed because I like graphics programs. They're fun, interesting, and produce pretty pictures! But I do feel guilty about specializing so much, so here's a more pythonic example.

You probably know that Python doesn't support multi-dimensional arrays. In fact, it doesn't support any arrays, it uses lists instead.

You can do something similar to a multi-dimensional array by making a list of lists. (Or list of lists of lists, etc.) It's not a true array but it's not so bad.

This example shows how you can use list comprehensions (or tuple, set, or dictionary comprehensions) to build these kinds of lists of lists.

Lists of Numbers

Suppose you want to build a list of lists where multiples[r] is a list holding multiples of the number r. Then multiples[r][c] holds the c-th multiple of r, which is r * c. This sort of acts like an array where r is the row number and c is the column number.

One way to build this kind of list of lists is to use nested loops as in the following code.

# Option 1: Build using two nested loops. multiples = [] for r in range(10): row = [] # Make this row's list. multiples.append(row) # Add it to the list of lists. for c in range(10): row.append(r * c)

Here the outer loop iterates over the "array's" rows and the inner loop fills in each row. Notice that cose uses r and c for the looping variables to help keep track of which is looping over rows and which is looping over columns.

This is a perfectly valid way to build the lists. It's effective and easy to understand so there's no shame in doing it this way.

Alternatively, you can use list comprehensions to build the lists. That's probably more pythonic but it's also more confusing, particularly when you're building lists of lists, so stick with the nested loops if that's easier to understand.

The trick I use to help me build lists of lists is to think about the innermost lists first. In this example, that means thinking about a "row" in the "array."

Here's some code that builds a single row.

# Step 1: Build an inner list multiples[J]. r = 3 # Just an example to get us started. row = [r * c for c in range(10)]

Behind the scenes, the list comprehension makes c loop through the columns 0 to 10 and appends the value r * c to the list.

If you run this code and print out the result, you get this:

[0, 3, 6, 9, 12, 15, 18, 21, 24, 27]

So far so good.

Now to make a list of those things, enclose that comprehension inside another one that loops over the rows. In the following code, I've highlighted the previous (inner) comprehension in blue.

multiples = [[r * c for c in range(10)] for r in range(10)]

Now for each row number r, the code creates a row of values and adds it to the list of lists. The following method prints the rows in a list of lists.

# Print the rows in a list of lists. def print_2d_list(list): for row in list: print(row) print()

Or you can use a comprehension.

print('\n'.join([f'{r}' for r in multiples]) + '\n')

If you print the multiple list's rows, you see this:

[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] [0, 2, 4, 6, 8, 10, 12, 14, 16, 18] [0, 3, 6, 9, 12, 15, 18, 21, 24, 27] [0, 4, 8, 12, 16, 20, 24, 28, 32, 36] [0, 5, 10, 15, 20, 25, 30, 35, 40, 45] [0, 6, 12, 18, 24, 30, 36, 42, 48, 54] [0, 7, 14, 21, 28, 35, 42, 49, 56, 63] [0, 8, 16, 24, 32, 40, 48, 56, 64, 72] [0, 9, 18, 27, 36, 45, 54, 63, 72, 81]

Huzzah! We have our list of multiples as desired.

Lists of Other Things

Hopefully at this point this seems so obvious that you're wondering why it's worth a post. With a little practice, you can probably build comprehensions like the previous one. I can usually do it in one or two tries. The situation becomes more confusing if you add in some objects, a set or dictionary, and some if clauses.

For example, suppose you have a list of Employee objects and you want to build a dictionary where departments['Shenanigans'] returns a list of the names of the employees in the Shenanigans department.

You can probably figure this sort of comprehension out. Go ahead and give it a try. I'll wait...

I can eventually get this sort of comprehension working, but it can be tricky trying to write it in one go. It's much easier if I follow my basic plan: start with the inner comprehensions and work my way out.

First, here's some code that defines the Employee class and makes a list of Employee objects.

# Another example with lists of eEmployee objects. # Define the Employee class. class Employee: def __init__(self, name, department): self.name = name self.department = department # Make some Employees. employees = [ Employee('Alice', 'Shenanigans'), Employee('Bob', 'Shenanigans'), Employee('Cindy', 'Mischief'), Employee('Dan', 'Mischief'), Employee('Eva', 'Knavery'), Employee('Frank', 'Mischief'), Employee('Gina', 'Knavery'), Employee('Horatio', 'Shenanigans'), ]

The goal is for departments['Shenanigans'] to return a list of employees names in the Shenanigans department, so let's start by trying to build a list of employees names in a particular department.

# Make departments{d} be a list of employees in department d. # Step 1a: Build an inner list departments[r]. d = 'Mischief' # "d" for "department." department_d = [employee.name for employee in employees if employee.department == d]

This comprehension selects employee names from the employees list. The if clause selects employees with the correct department, which is represented by the variable d. If you run this code and print the result, you'll get the following.

['Cindy', 'Dan', 'Frank']

Next, we need to make variable d loop through the department names. Unfortunately, we don't have a handy list of department names. Fortunately, we can use a comprehension to build one.

# Step 1b: Build a set of department names. department_names = {employee.department for employee in employees}

This comprehension loops through the employees and adds their department names to a set. If you're unfamiliar with sets, they're kind if like lists that do not allow duplicates. If you add the string "apple" to a set twice, the set only keeps one copy of "apple."

You can tell that this comprehension builds a set because (1) it uses squiggly brackets { and }, and (2) it doesn't use colons : the way you would to create a dictionary.

Now we can move on to step 2 where we wrap the first comprehension inside another one. This time we make the variable d from the previous comprehension loop over the department names. The following code shows the new version with the original comprehension highlighted in blue.

# Step 2: Wrap that comprehension in another one, # replacing d with a looping variable. departments = { d: [employee.name for employee in employees if employee.department == d] for d in department_names}

This comprehension makes d loop through the department names. For each department, it makes variable employee loop through the employee list and it picks out the employees with the right department name. It then saves d and the employee's name in a dictionary. (You can tell it's a dictionary because it uses both squiggly brackets and colons.)

After running this comprehension, the example program uses the following code to display the resulting dictionary's keys and values.

Knavery : ['Eva', 'Gina'] Mischief : ['Cindy', 'Dan', 'Frank'] Shenanigans: ['Alice', 'Bob', 'Horatio']

This is just what we want! At this point we could call mom to brag about our fantasic accomplishment, or we could make the code even more pythonic by piling all of the code into One Comprehension to Rule Them All.

# Step 3: Replace department_names with its own comprehension. departments = { d: [employee.name for employee in employees if employee.department == d] for d in {employee.department for employee in employees}}

This gives the same result as the previous version.

Now everything is all wrapped up in a single fairly confusing comprehension. It's actually not too hard to understand if you read it slowly, but building it in the first place can be rather tricky.

Conclusion

Building these kinds of comprehensions is easier if you stick to a basic plan. Build the innermost comprehension first using a placeholder variable like r in the first example and d in the second. Test it and, when it's working, embed it in another comprehension that makes the placeholder variable loop over something. Continue assembling the pieces until you have a single grand unified comprehension.

These kinds of deep comprehensions can be tricky to build and confusing to read. If you prefer to use a series of nested loops, by all means do so. I often include both versions (with one commented out) mostly to see if I can figure it out how to do it. The looping version is sometimes the better option if for no other reason than it makes long-term maintenance easier, but sometimes the comprehensions make a fun challenge. (And they're great conversation starters at parties.)

Download the example to experiment with it and to see additional details.