Python tricks and techniques for concise Leetcode
During a recent period of unemployment, I fell in with a crowd feeding my worst addictions. This group and I would get together most days and work on Leetcode problems. Even when I told myself I would stop, I kept coming back for one more hit of elegant solutions to short, focused problems.
I learned some tricks and new Python features from this group, and hope you, dear reader, will find them useful. Features are ordered roughly by complexity- scroll down until you find things that seem interesting. If you hit the end of the page before finding something interesting, write your own blog post and send it to me.
- Leetcode auto-imports lots of stuff
- Product is a great way to merge for loops
- min, max, sum, prod, any, all, next
- Reversing with [::-1]
- The walrus operator
- lru_cache
- heapq- largest/smallest k
- String operations are really fast in Python
- BYOIterators
- Subtract counters to use as a boolean
- Conclusion
Leetcode auto-imports lots of stuff
The Python3 runtime includes, among other things, collections, functools, heapq, itertools, bisect, statistics, and re. That’s a pretty good hint these things will be useful. Importing everything into one namespace probably isn’t great practice generally, but it makes some code much cleaner- consider the difference between
return sum(math.floor(n/(5**i))
for i in range(1,math.ceil(math.log(n,5))+1)
) if n > 4 else 0
vs
return sum(floor(n/(5**i))
for i in range(1,ceil(log(n,5))+1)
) if n > 4 else 0
So, go ahead and take advantage of that. Lots of the following code sure does.
Product is a great way to merge for loops
itertools is a module that’s filled with nifty tools for leetcode, but the one I used most often, and maybe most inappropriately, was itertools.product (auto-imported as product). It generates all combinations of iterables, and returns them as tuples that can be auto-unpacked. In addition to its intended use in Cartesian products, it’s a great way to cut those multi-line for loops down to one line. For example, in the Toeplitz Matrix problem, you could use1:
class Solution:
def isToeplitzMatrix(self, matrix: List[List[int]]) -> bool:
for i in range(len(matrix)):
for j in range(len(matrix[0])):
if matrix[i][j] != matrix[max(0,i-j)][max(0,j-i)]:
return False
return True
or
class Solution:
def isToeplitzMatrix(self, matrix: List[List[int]]) -> bool:
for i, j in product(range(len(matrix)), range(len(matrix[0]))):
if matrix[i][j] != matrix[max(0,i-j)][max(0,j-i)]:
return False
return True
Product also has a nifty repeat parameter, so instead of:
for i, j in product(range(len(nums)), range(len(nums))):
# Do something
you could say:
for i, j in product(range(len(nums)), repeat=2):
# Do something
saving single digits of characters, at the mere cost of making readers of your code need to deeply understand the itertools product function.
This also works well in list/generator comprehensions, so you could solve the Toeplitz Matrix problem with:
return all(matrix[i][j] == matrix[max(0,i-j)][max(0,j-i)]
for i, j in product(range(len(matrix)),
range(len(matrix[0]))))
Speaking of,
min, max, sum, prod, any, all, next
These are my favorite ways to turn an iterable into a single value, thereby stretching list and generator expressions to the limit. min, max, sum, and prod all do what they sound like- prod being the relatively new and controversial “Multiply all the elements together” function. any and all take in iterators to things that can be treated as booleans, and return whether any of those or all of those resolve to True-ish respectively. Keep in mind that they don’t need to be booleans exactly- 0, None, empty lists, etc. are just as good as False, and lots of “Normal” objects and values will be treated as True. next() was particularly useful when I had a generator expression that I was sure would resolve something, and just wanted the first value- e.g. from the intersection of two linked lists
class Solution:
def getIntersectionNode(self, headA: ListNode, headB: ListNode) -> ListNode:
hA = set(headA)
return next((b for b in headB if b in hA), None)
(The astute reader will note that you can’t call set() or iterate on a listNode- more on this later)
All of these functions complain when given an empty list/generator expression with every element filtered, but most of them allow for a default value so you don’t need to handle the empty case elsewhere. This is especially nice when Leetcode tells you explicitly what to return in these semi-error cases, like in the linked list intersection case.
Reversing with [::-1]
Sure, you could reverse a list with “reversed(l)”, but that’s so many wasted characters compared to indexing each element of the array backward! Plus, reversed(l) returns an iterator instead of a real list, which is sometimes great, but other times verbose. Consider this solution to product of array except self
class Solution:
def productExceptSelf(self, nums: List[int]) -> List[int]:
def prefix_op(nums, f, base):
prev = base
result = []
for num in nums:
prev = f(prev, num)
result.append(prev)
return result
forward = prefix_op(nums, lambda x, y: x*y, 1)
backward = list(reversed(prefix_op(reversed(nums), lambda x, y: x*y, 1)))
return backward[1:2]\
+ [forward[i-1] * backward[i+1] for i in range(1, len(nums)-1)]\
+ forward[-2:-1]
Does the “backward” line look better as
backward = list(reversed(prefix_op(reversed(nums), lambda x, y: x*y, 1)))
or
backward = prefix_op(nums[::-1], lambda x, y: x*y, 1)[::-1]
Tell me the second version doesn’t feel more like the matrix.
The walrus operator
The operator is a relatively new (python 3.8) feature that merges evaluating a value and assigning it- no more
v = expression
if v
Now you can just say:
if v := expression
and v will be good to go!
It’s also nice for some list comprehensions- the prefix_op() function above could be rewritten
def prefix_op(nums, f, base):
prev = base
return [prev:=f(prev, num) for num in nums]
lru_cache
Lot of problems in leetcode are most efficient as dynamic programming problems, where you build up a solution from smaller problems. Traditionally, you compute the solutions for all subproblems, then use those solutions to solve a larger problem, in a bottom-up approach. But it’s often easier to think about problems from the top down, and form the solution recursively.
“Traditional” dynamic programming solution to the unique paths problem:
class Solution:
def uniquePaths(self, m: int, n: int) -> int:
unique_paths = [[0] * (n+1)] * (m+1)
unique_paths[1][1] = 1
for x in range(1,m+1):
for y in range(1,n+1):
if x == 1 and y == 1:
continue
unique_paths[x][y] = unique_paths[x-1][y] + unique_paths[x][y-1]
return unique_paths[m][n]
The recursive solution:
class Solution:
@lru_cache(None)
def uniquePaths(self, m: int, n: int) -> int:
if m <= 1 or n <= 1:
return 1
return self.uniquePaths(m-1, n) + self.uniquePaths(m, n-1)
If you throw in a @lru_cache(None) over a function, it will memoize the results of each function call to avoid recomputation- in practice, this tends to mean the same things get calculated and stored in RAM as would be in the dynamic program, but it might be much cleaner.
The None in the @lru_cache() specifies the limit on how many entries to cache- in leetcode, None (no limit) is often fine. If you’re being really RAM-conscious, think through how many values really need to be remembered. For example, the recursive fibonacci sequence only really needs to remember the most recent 3 entries- the value being currently calculated, and the two previous values.
@lru_cache(3)
def fibonacci(n):
return n < 2 ? n : fibonacci(n-1) + fibonacci(n-2)
Be a little careful with the @lru_cache- it avoids recomputation, but if you have a call stack that’s very deep, you’ll still get a stack overflow and will need to write a non-recursive version, or do something especially clever to preserve the recursion.
heapq- largest/smallest k
A decent chunk of leetcode problems ask for K values best matching some criteria- the heapq functions nlargest and nsmallest make this relatively straightforward. For example, to find the most frequent elements, you can count the elements, find the K highest element counts and their associated elements, and return those.
def topKFrequent(self, nums: List[int], k: int) -> List[int]:
c = Counter(nums)
# Sort by count, keep k largest
max_keys = nlargest(k, ((v, i) for i, v in c.items()))
return list(zip(*max_keys))[1] # Get the indices
Before you think “Why is that a function, I could just X”, check the benchmarks on nlargest/nsmallest- it’s really fast.
String operations are really fast in Python
If I’ve learned anything from over-optimizing leetcode submissions, it’s that the timer on leetcode is super noisy and not suitable for real benchmarking.
But if I’ve learned two things, it’s that string operations are really fast in Python. For example, to convert a list of numbers to representing decimal digits to an int,
int("".join(str(d) for d in digs),10)
is about 10x as fast as
sum(v * 10 **i for i, v in enumerate(digs))
(whether we reverse the digits or not). If you were storing digits as strings to begin with,
int("".join(digs_s),10)
is about 10x faster than that.
This is the opposite of my intuition in c++, where I’ve been told to try hard to avoid strings. I think the key is that string operations are implemented in C, so the interpreter is needed less?
BYOIterators
Leetcode constantly uses linked lists as a data structure for problems that would be much easier with standard lists. This means you lose all kinds of convenience functions- you can’t call set() or Counter() on a linked list node, because list nodes aren’t iterable. But, since this is Python, feel free to just add an iteration capability2.
def LinkedListIterator(self):
n = self
while n:
yield n
n = n.next
ListNode.__iter__ = LinkedListIterator
This is the trick that made the intersection node problem above work- the full solution is
def LinkedListIterator(self):
n = self
while n:
yield n
n = n.next
ListNode.__iter__ = LinkedListIterator
class Solution:
def getIntersectionNode(self, headA: ListNode, headB: ListNode) -> ListNode:
hA = set(headA)
return next((b for b in headB if b in hA), None)
The “main” function is now much easer- iteration allows us to put all of A in a set in one call, and trivially iterate through nodes of B.
Subtract counters to use as a boolean
You can subtract one counter from another (A-B), returning the number of times each element occured in A minus the number of times it occured in B. What if an element occured more in B than A? Well, counters should’t have negative entries, so those elements are just excluded. That gives a surpisingly nice, if inefficient, way to tell if one collection is a superset of another- just subtract the counters, and treat the resulting difference as a boolean, since an empty counter will resolve as False. For example
def minWindow(self, s: str, t: str) -> str:
start = stop = 0
windowContains, needs = Counter(), Counter(t)
res = None
while stop != len(s) or start != len(s):
if needs - windowContains and stop < len(s):
windowContains[s[stop]] += 1
stop += 1
else:
windowContains[s[start]] -= 1
start += 1
if not (needs - windowContains):
res = res if res and len(res) < stop-start else s[start:stop]
return res if res else ""
The line
if needs - windowContains
checks whether every character in the target string occurs at least as many times in the string window we’re looking at.
Conclusion
Do these make your code cleaner, clearer, or a better fit for production? No idea, see the name of this site. But they fit fantastically on a whiteboard.
Special thanks to Matthew Caseres, Tony Chow, Santiago Gonzalez, Andrew Lee, and many others for the work put into Leetcode practice.