enumerate and zip
enumerate() gives you the index along with each item. zip() lets you walk through two or more lists side by side. Here's the short version of both:
Both show up constantly in interview code. Once you're used to them, your loops get shorter and a lot easier for the interviewer to follow.
enumerate(): index and value together
O(1) extraenumerate() hands you a pair on every pass: the index and the item. The for i, x in ... part unpacks that pair into two variables:
By default it counts from 0. Pass start if you want something else, like numbering lines for a person to read:
Note
start only changes the number you get back. It doesn't skip any items. enumerate(steps, start=1) still begins at "wash", so don't use that number to index back into the list.
Why enumerate beats range(len(...))
A lot of people come to Python from languages where you loop over indices, so they write this:
It works, but you have to read nums[i] to figure out what the loop is about. With enumerate(), the value has a name right in the loop header:
Here's a real example. In Two Sum, you need both the number and where it was, and enumerate() gives you both:
Note
range(len(...)) still makes sense when you don't need the value at all, or when you're jumping around by index, like looking at nums[i + 1] or walking two pointers. Use enumerate() when you'd otherwise write nums[i] on the first line of the loop.
zip(): loop over lists in parallel
O(1) extrazip() takes the first item from each list, then the second from each, and so on. Each step gives you a tuple:
You can combine it with enumerate() when you need the index too. Just wrap the zip() and unpack the inner tuple with parentheses:
Note
zip() gives you an iterator, not a list. That's why printing it directly shows something like <zip object>. Wrap it in list() to see the pairs. It also means you can only loop over it once.
Transposing a matrix with zip(*matrix)
O(rows·cols)The * unpacks the matrix, so each row becomes its own argument to zip(). Then zip() groups the first item of every row, then the second, and so on. Those groups are the columns:
This comes up in matrix problems. In Rotate Image, for example, turning a matrix 90 degrees clockwise is the same as reversing the rows and then transposing:
Building a dict with dict(zip(keys, values))
If you have one list of keys and one list of values, zip them and pass the result to dict():
A similar trick maps each item to its index. This one uses enumerate() inside a dict comprehension:
Pairs of neighbors with zip(nums, nums[1:])
Zipping a list with itself shifted by one gives you every pair of neighbors. It's a clean way to compare each item with the next one:
Note
nums[1:] makes a copy of the list, so this uses O(n) extra space. That's fine almost all the time. If an interviewer asks for O(1) space, loop with range(1, len(nums)) and compare nums[i - 1] to nums[i] instead.
zip() stops at the shortest list
watch outIf the lists aren't the same length, zip() quietly stops when the shortest one runs out. The leftover items just disappear, and you don't get an error:
You probably won't need this in an interview, but it's handy to know. If you want to keep the extras, use itertools.zip_longest. It fills the gaps with None, or with whatever you pass as fillvalue:
The itertools page covers zip_longest and the rest of that module. There's also zip(a, b, strict=True) on Python 3.10+, which raises an error when the lengths don't match, but that's more of a production-code thing than an interview one.