The Quest Begins (The "Why")
I still remember my first regional coding contest. I was pumped, my laptop humming like a DeLorean, and I opened the problem set with a grin. Thirty minutes in, I hit a wall: a problem that asked for the k smallest sums from two sorted arrays. My naïve solution generated all possible pairs, sorted them, and then sliced the first k. It worked on the sample, but the judge returned Time Limit Exceeded with a sneaky smile. I felt like I’d just watched the hero get knocked out in the first round of a tournament—frustrating, embarrassing, and utterly motivating.
That night I dug into forums, watched a few livestreams, and kept stumbling upon snippets that made veterans nod approvingly. “Why didn’t I know this?” I thought. The answer wasn’t a new algorithm; it was a handful of Python tricks that most tutorials skip because they seem “too minor.” Yet, in the pressure cooker of competitive programming, those tiny shortcuts shave off precious milliseconds—and sometimes they’re the difference between a green check and a red X.
Let’s grab your own quest, and I’ll share the three features that turned my TLEs into ACs.
The Revelation (The Insight)
1. heapq.nlargest / nsmallest – Get the top‑k without sorting everything
Most of us reach for sorted(arr)[:k] when we need the k smallest (or largest) items. It’s clean, but it’s O(n log n) because we sort the whole list even if we only need a handful of elements. The heapq module gives us nlargest and nsmallest, which run in O(n log k)—a huge win when k ≪ n.
The gotcha? The functions return a list already sorted in the requested order, but they don’t preserve the original order of equal elements. If stability matters, you’ll need to add a tie‑breaker (like the original index) to the key.
2. bisect – Binary search on a sorted list in pure Python
When you need to insert into a sorted container or find the first element ≥ x, writing your own binary‑search loop is error‑prone (off‑by‑one bugs love to hide in the while lo < hi: condition). The built‑in bisect module does it reliably in C‑speed, exposing bisect_left, bisect_right, and insort*.
The trap? insort still shifts elements, so it’s O(n). It’s perfect when you’re doing a few inserts and many queries, but if you need lots of dynamic updates, a balanced tree (or sortedcontainers library) would be better.
3. itertools.accumulate – Running totals, prefixes, and more with one line
Prefix sums are a staple in CP: you want to know the sum of a sub‑array in O(1) after O(n) preprocessing. The classic way is a manual loop that appends to a list. itertools.accumulate does the same thing, but you can also plug in any binary function—think running product, running max, or even a custom modulo sum.
The nuance? accumulate returns an iterator, not a list. If you forget to wrap it with list() (or iterate it only once), you’ll get surprising empty results later.
Wielding the Power (Code & Examples)
Problem: K Smallest Pair Sums
Given two sorted arrays A and B, find the k smallest sums A[i] + B[j].
The Struggle (Before)
def k_smallest_sums_bruteforce(A, B, k):
sums = [a + b for a in A for b in B] # O(n*m) memory & time
sums.sort()
return sums[:k]
If A and B each have 10⁴ elements, we create 10⁸ sums—far too much.
The Victory (After) – using a min‑heap + heapq.heappushpop
import heapq
def k_smallest_sums_heap(A, B, k):
# We only need to explore the first k elements of each array.
# Push (sum, i, j) where i, j are indices.
heap = [(A[0] + B[0], 0, 0)]
visited = {(0, 0)}
result = []
while heap and len(result) < k:
s, i, j = heapq.heappop(heap)
result.append(s)
# Next candidate: move in A
if i + 1 < len(A) and (i + 1, j) not in visited:
heapq.heappush(heap, (A[i + 1] + B[j], i + 1, j))
visited.add((i + 1, j))
# Next candidate: move in B
if j + 1 < len(B) and (i, j + 1) not in visited:
heapq.heappush(heap, (A[i] + B[j + 1], i, j + 1))
visited.add((i, j + 1))
return result
Why this shines: we never generate more than O(k log k) entries, and the heap stays tiny. The heapq functions are implemented in C, so the constant factor is low.
Problem: Count Frequencies & Retrieve Top‑M
You have a list of integers and need the M most frequent values.
The Struggle (Before)
def top_m_bruteforce(nums, m):
freq = {}
for x in nums:
freq[x] = freq.get(x, 0) + 1 # manual get
# sort by frequency descending
sorted_items = sorted(freq.items(), key=lambda kv: kv[1], reverse=True)
return [val for val, _ in sorted_items[:m]]
Sorting the whole dictionary is O(u log u) where u is the number of unique values—often close to n.
The Victory (After) – collections.Counter + heapq.nlargest
from collections import Counter
import heapq
def top_m_heap(nums, m):
cnt = Counter(nums) # O(n) counting
# nlargest works directly on the items; key extracts the count
most_common = heapq.nlargest(m, cnt.items(), key=lambda kv: kv[1])
return [val for val, _ in most_common]
Counter is a subclass of dict optimized for counting, and heapq.nlargest avoids the full sort. The gotcha: if two numbers have the same frequency, the order returned by nlargest is undefined (it follows the heap’s internal ordering). Add a secondary key (e.g., the number itself) if you need deterministic output.
Problem: Prefix Sums with Modulo
You need to answer many queries of the form “sum of A[l..r] modulo P”.
The Struggle (Before)
def prefix_mod_bruteforce(A, P):
pref = [0]
for x in A:
pref.append((pref[-1] + x) % P) # manual loop
return pref
Fine, but you repeat the same pattern for product, max, etc.
The Victory (After) – itertools.accumulate
from itertools import accumulate
import operator
def prefix_mod_accumulate(A, P):
# accumulate returns an iterator; we cast to list for random access
return list(accumulate(A, lambda acc, x: (acc + x) % P, initial=0))
Now change the lambda to operator.mul for a running product, or max for a running maximum—all with the same one‑liner. Remember: if you later do pref = accumulate(...) and then try to index pref[5] twice, the iterator will be exhausted on the first use. Wrap it in list() (or reuse the iterator only once) to avoid that pitfall.
Why This New Power Matters
These three tricks taught me that competitive programming isn’t just about knowing the right algorithm—it’s about wielding the language’s standard library like a seasoned swordsman.
-
Speed:
heapq.nlargest/nsmallestandbisectcut the asymptotic complexity of common subtasks without writing extra code. -
Clarity: A single line with
accumulatereplaces a boiler‑plate loop, making the intent instantly readable. - Robustness: Leveraging battle‑tested modules reduces the chance of off‑by‑one bugs that haunt manual implementations.
When I started using them, my submission times dropped from seconds to milliseconds, and my confidence soared. I stopped fearing the “TL;DR” label and started looking forward to the next challenge, knowing I had a few hidden spells in my grimoire.
Your Turn – A Small Quest
Pick a problem you’ve solved before with a naïve O(n²) or O(n log n) approach. Refactor it using one of the tricks above (try heapq.nlargest for a top‑k filter, bisect for a search/insert, or accumulate for a running aggregate). Share your before/after snippets in the comments—let’s see who can shave off the most milliseconds!
Happy coding, and may your bugs be few and your ACs many! 🚀
Top comments (0)