Udaiy’s Blog

Python and Tensor Notes

Part I — Python

Complexity at a glance

n up to acceptable typical tool
10⁶+ O(n) / O(n log n) one pass, dict/set, sort
~10⁴ O(n²) nested loops fine
~20 O(2ⁿ) brute-force subsets / recursion fine

x in set/dict is O(1); x in list is O(n). Prefer deque for front inserts/pops.

Containers

xs.append(v); xs.pop()                 # O(1) at the end
xs[::-1]; xs[1:4]; xs[::2]            # slices copy
d.get(k, default); d.setdefault(k, []).append(v)
max(d, key=d.get)                      # key with largest value
seen = set(); a & b; a | b; a - b
d[(row, col)] = v                      # tuple keys

Strings

s.lower(); s.strip(); s.split(); s.split(",")
"-".join(parts)                        # parts must be strings
s.find("sub")                          # -1 if missing
ord("a"); chr(97); ord(c) - ord("a")   # letter → 0..25
s == s[::-1]                           # palindrome
"".join(sorted(s))                     # anagram key

Build with a list + "".join, not s += ch in a hot loop.

Iteration

for i, v in enumerate(xs): ...
for a, b in zip(xs, ys): ...
list(zip(*matrix))                     # transpose
any(...); all(...)
{v: i for i, v in enumerate(xs)}       # value → index
sorted(xs, key=lambda p: (-p[1], p[0]))
max(range(len(xs)), key=xs.__getitem__)  # argmax
a, b = b, a

Standard library that pays rent

from collections import Counter, defaultdict, deque

Counter(xs).most_common(2)
Counter(a) == Counter(b)               # anagrams
d = defaultdict(list); d[k].append(v)
q = deque([start]); q.append(v); q.popleft()

import heapq
heapq.heappush(h, v); heapq.heappop(h)
heapq.nlargest(3, xs)

import bisect
bisect.bisect_left(xs, v)              # xs must be sorted

from itertools import combinations, permutations, product, accumulate
from functools import lru_cache

@lru_cache(maxsize=None)
def fib(n):
    return n if n < 2 else fib(n - 1) + fib(n - 2)

import math, random
math.gcd(a, b); math.lcm(a, b); math.comb(n, k); math.isqrt(n)
math.inf                               # seed for min-scans
random.randint(a, b)                   # BOTH ends inclusive
random.choice(xs); random.sample(xs, k)

Pattern templates

Two-sum (hash map)

def two_sum(nums, target):
    seen = {}
    for i, v in enumerate(nums):
        if target - v in seen:
            return [seen[target - v], i]
        seen[v] = i

Two pointers

lo, hi = 0, len(xs) - 1
while lo < hi:
    s = xs[lo] + xs[hi]
    if s == target: return lo, hi
    if s < target: lo += 1
    else: hi -= 1

Sliding window (longest unique substring)

def longest_unique(s):
    last, left, best = {}, 0, 0
    for right, ch in enumerate(s):
        if ch in last and last[ch] >= left:
            left = last[ch] + 1
        last[ch] = right
        best = max(best, right - left + 1)
    return best

Running best (max profit / Kadane)

def max_profit(prices):
    best, low = 0, math.inf
    for p in prices:
        low = min(low, p)
        best = max(best, p - low)
    return best

def max_subarray(nums):
    best = cur = nums[0]
    for v in nums[1:]:
        cur = max(v, cur + v)
        best = max(best, cur)
    return best

Stack (valid parentheses)

def valid_parentheses(s):
    pairs = {")": "(", "]": "[", "}": "{"}
    stack = []
    for ch in s:
        if ch in "([{":
            stack.append(ch)
        elif not stack or stack.pop() != pairs[ch]:
            return False
    return not stack

Prefix sums

prefix = [0]
for v in nums:
    prefix.append(prefix[-1] + v)
# sum(nums[i:j]) == prefix[j] - prefix[i]

Binary search

def binary_search(xs, target):
    lo, hi = 0, len(xs) - 1
    while lo <= hi:
        mid = (lo + hi) // 2
        if xs[mid] == target: return mid
        if xs[mid] < target: lo = mid + 1
        else: hi = mid - 1
    return -1  # or return lo for insertion index

BFS on a grid

from collections import deque

def bfs(grid, start):
    rows, cols = len(grid), len(grid[0])
    q = deque([(*start, 0)])
    seen = {start}
    while q:
        r, c, dist = q.popleft()
        for dr, dc in ((1, 0), (-1, 0), (0, 1), (0, -1)):
            nr, nc = r + dr, c + dc
            if 0 <= nr < rows and 0 <= nc < cols and (nr, nc) not in seen \
               and grid[nr][nc] != "#":
                seen.add((nr, nc))
                q.append((nr, nc, dist + 1))

1-D DP / merge intervals

def climb_stairs(n):
    a, b = 1, 1
    for _ in range(n - 1):
        a, b = b, a + b
    return b

def merge_intervals(intervals):
    out = []
    for s, e in sorted(intervals):
        if out and s <= out[-1][1]:
            out[-1][1] = max(out[-1][1], e)
        else:
            out.append([s, e])
    return out

Number quirks

divmod(17, 5)          # (3, 2)
-7 // 2                # -4  (floors toward -∞)
n % 10; n // 10        # peel / drop last digit
round(2.5)             # 2  (banker's rounding)

Gotchas

  1. def f(xs=[]) keeps state across calls — use xs=None.
  2. [[0]*3]*3 aliases one row — use [[0]*3 for _ in range(3)].
  3. b = a aliases; a[:] / a.copy() are shallow copies.
  4. Don't mutate a list while iterating it.
  5. is ≠ == for values; use == for numbers and strings.
  6. range(a, b) excludes b; random.randint(a, b) includes both.
  7. sorted("ab") returns a list of chars — "".join it.
  8. xs.sort() returns None.
  9. Recursion depth ≈ 1000 — prefer iteration for deep chains.

Pre-submit checklist

Empty / one-element / all-same / negatives / already sorted / off-by-one at ends / exact return type (-1 vs None, list vs tuple).

Part II — Tensors (PyTorch + einops)

Always ask three things

x.shape; x.dtype; x.device; x.numel()

t.tensor([1, 2]) is int64; call .float() before mean / norm / std. Use .item() when a plain Python number is required.

Creating & reshaping

t.arange(3, 9)             # end exclusive
t.linspace(0, 1, 5)        # end inclusive
t.zeros(2, 3); t.ones_like(x); t.rand(2, 3); t.randn(2, 3)
x.reshape(a, b)            # prefer over view
x.T; x.permute(2, 0, 1); x.squeeze(d); x.unsqueeze(d)

Indexing

x[0]; x[:, 1]; x[1, 2]; x[..., 0]; x[:, None]
x[x > 7]                   # boolean mask → 1D
prices[items]              # fancy index; shape follows `items`
mat[rows, cols]
mat[tuple(coords.T)]       # coords: (batch, ndim)
matrix[row_indexes]        # pick rows
matrix[:, col_indexes]     # pick columns

Broadcasting (align from the right)

  1. Prepend size-1 dims until ranks match.
  2. Size-1 dims stretch; any other mismatch errors.
(3, 1, 5) + (1, 4, 5) → (3, 4, 5)   OK
(8, 2, 6) + (2, 6)    → (8, 2, 6)   OK
(4, 1) + (4,)         → (4, 4)      silent trap

Use keepdim=True after reductions that must broadcast back:

x / x.sum(dim=1, keepdim=True)

Reductions

dim is the axis that disappears.

x.sum(dim=0)               # column sums
x.mean(dim=1, keepdim=True)
x.max(dim=1).values        # .max(dim=...) returns (values, indices)
x.argmax(dim=1); x.norm(dim=1)

Note: torch.std is unbiased (÷ n−1) by default; NumPy's std is biased (÷ n).

einops

Parentheses (a b): first name varies slowest. (2 h) = block copies; (h 2) = elementwise stretch.

einops.rearrange(x, "b c h w -> b (c h w)")
einops.rearrange(x, "(b1 b2) c h w -> c (b1 h) (b2 w)", b1=2)
einops.rearrange(t.arange(3, 9), "(h w) -> h w", h=3, w=2)

einops.repeat(v, "w -> (w 7)")         # each value repeated 7×
einops.repeat(v, "w -> (7 w)")         # whole vector tiled 7×

einops.reduce(temps, "(w 7) -> w", "mean")
einops.reduce(x, "b c (h 2) (w 2) -> b c h w", "max")  # 2×2 max-pool

Group z-score without loops:

avg = einops.reduce(temps, "(w 7) -> w", "mean")
std = einops.reduce(temps, "(w 7) -> w", t.std)
out = (temps - einops.repeat(avg, "w -> (w 7)")) / einops.repeat(std, "w -> (w 7)")

einsum (three rules)

  1. Name repeated across inputs → multiply.
  2. Name missing from output → sum away.
  3. Output names may be reordered freely.
einops.einsum(A, "i i ->")                    # trace
einops.einsum(A, v, "i j, j -> i")             # matvec
einops.einsum(A, B, "i j, j k -> i k")         # matmul
einops.einsum(u, v, "i, i ->")                 # dot
einops.einsum(u, v, "i, j -> i j")             # outer
einops.einsum(X, Y, "b i j, b j k -> b i k")   # batched matmul
einops.einsum(Q, K, "b q d, b k d -> b q k")   # attention scores

gather

Output shape equals index shape.

# dim=1: out[i, j] = input[i, index[i, j]]
matrix.gather(1, indexes)

# pick class score per example
logits.gather(1, labels.unsqueeze(1)).squeeze(1)
# same:
logits[t.arange(len(labels)), labels]

Softmax family (numerically stable)

Subtract the row max so exp never overflows.

def softmax(x):
    x = x - x.max(dim=-1, keepdim=True).values
    e = x.exp()
    return e / e.sum(dim=-1, keepdim=True)

def logsumexp(x):
    C = x.max(dim=-1).values
    return C + (x - C.unsqueeze(-1)).exp().sum(dim=-1).log()

def log_softmax(x):
    return x - logsumexp(x).unsqueeze(-1)

def cross_entropy(logits, labels):
    logprobs = log_softmax(logits)
    return -logprobs.gather(1, labels.unsqueeze(1)).squeeze(1)

Chain: logsumexp = max + log Σ exp(x − max) → log_softmax = x − logsumexp → CE = −log_softmax[label].

Sampling & accuracy

def sample_distribution(probs, n):
    return (t.rand(n, 1) > t.cumsum(probs, dim=0)).sum(dim=-1)

(scores.argmax(dim=1) == true_classes).float().mean()

Autograd loop

w = t.randn(1, requires_grad=True)
b = t.zeros(1, requires_grad=True)

for _ in range(200):
    loss = ((w * x + b - y) ** 2).mean()
    loss.backward()
    with t.no_grad():
        w -= lr * w.grad
        b -= lr * b.grad
    w.grad.zero_(); b.grad.zero_()

Chant: zero → forward → loss → backward → step. Updates live inside no_grad.

Linear algebra one-liners

Mn = matrix / matrix.norm(dim=1, keepdim=True)   # L2-normalise rows
Mn @ Mn.T                                        # pairwise cosine similarity
A @ B; u @ v; t.outer(u, v); t.trace(A)

Composition of linear maps is linear — networks need nonlinearities or depth collapses to one matrix.

Attention (scaled dot-product)

scores = einops.einsum(Q, K, "b q d, b k d -> b q k") / math.sqrt(Q.shape[-1])
# optional causal mask: scores.masked_fill(triu_mask, float("-inf"))
out = softmax(scores) @ V

/sqrt(d_k) keeps softmax out of the saturated regime when dimension grows.

Tensor bug checklist

  1. Output shape matches the docstring ((n,) ≠ (n, 1)).
  2. .float() before stats; .long() for indices; .item() for Python scalars.
  3. keepdim=True when broadcasting after a reduce.
  4. max(dim=...) returns a tuple — take .values.
  5. Cross-entropy / BCE-with-logits want logits, not probabilities.
  6. Silent (n,1)+(n,)→(n,n) broadcast when something looks “squared.”
  7. In (a b), first name varies slowest.

Habit that saves the most time

Before writing code, jot the contract: inputs → outputs (types or shapes) and one tiny example. Most wrong answers are misread specs or shape bugs, not missing algorithms.