Avoid nested for loops with itertools.product() (python)
The Python itertools module contains many useful iterator functions. This entry will focus on the combinatoric iterators.
Today I solved the advent of code puzzle using nested for loops to iterate over each item in an array twice:
for i in data_list:
for j in data_list:
if i + j == 2020:
print(i * j)
This nested for loop is finding the “cartesian product” of the the arrays it is looping over. It finds the set of all ordered pairs that can be constructed from the two arrays. The operation can be done more efficiently and succinctly by using ‘itertools.product()’.
for n in itertools.product(data_list, data_list):
if sum(n) == 2020:
print(np.product(n))
The itertools module also has functions for combinations and permutations as well:
Iterator |
Arguments |
Results |
|---|---|---|
p, q, … [repeat=1] |
cartesian product, equivalent to a nested for-loop |
|
p[, r] |
r-length tuples, all possible orderings, no repeated elements |
|
p, r |
r-length tuples, in sorted order, no repeated elements |
|
p, r |
r-length tuples, in sorted order, with repeated elements |
Examples |
Results |
|---|---|
|
|
|
|
|
|
|
|