Vibe Engines
YouTube

Vibe Engines is free, interactive AI engineering & system design — runs in your browser. Pick a learning path

Pandas & NumPy

The data-wrangling operations every ML notebook runs on — selection, groupby, joins, reshaping, vectorization.

NumPy provides the ndarray: a typed block of numbers on which every operation runs in compiled code over the whole array at once. pandas builds the DataFrame on top of it: labeled columns, each with its own dtype, plus an index that aligns rows in every operation.

The mental model is the same for both. Say what to do to a whole column, never loop over rows in Python, and let alignment and broadcasting line the data up. The rows cover pandas 2 and 3; pandas 3 turns on Copy-on-Write and a dedicated string dtype by default, and the gotchas that changes are marked.

Load & inspect

pd.read_csv(path)
set usecols=, dtype= and parse_dates= at load time instead of fixing columns later
chunksize=100_000
with read_csv: stream a file too big for memory as an iterator of frames
pd.read_parquet(path)
columnar, typed and compressed — far faster than CSV; needs pyarrow
df.info() · df.describe()
dtypes, non-null counts and memory · summary statistics per numeric column
df.head() · df.sample(5)
first rows · random rows — a sample shows more than the top of a sorted file
s.value_counts()
frequency of each value; normalize=True for shares, dropna=False to count NaN
df.nunique()
distinct values per column — spot IDs, constants and category candidates
to_csv(p, index=False)
write without the index, which otherwise comes back as an Unnamed: 0 column

Selection (the part everyone fumbles)

df.loc[rows, cols]
by LABEL / boolean mask — df.loc[df.x > 5, "y"]
df.iloc[i, j]
by POSITION — iloc[0] first row, iloc[:, -1] last col
Boolean masks
combine with & | ~ and parentheses — never and/or
df.query("x > 5")
readable chained filters, e.g. "x > 5 and y == 'a'"; @var reads a Python variable
SettingWithCopyWarning
assign with one .loc[mask, col] = val; chained df[mask][col] = val never writes in pandas 3
df[["a", "b"]]
double brackets return a DataFrame; single df["a"] returns a Series
s.isin(["a", "b"])
membership mask, the vectorized in; ~ negates it
s.between(1, 10)
range mask, inclusive at both ends; inclusive="left" makes it half-open

Clean & transform

df.isna().sum()
the null audit; then fillna / dropna(subset=…)
drop_duplicates(subset)
dedupe on key columns; keep="last" keeps the newest instead of the first
df.assign(c=lambda d: …)
add columns inside a method chain; the lambda sees the frame built so far
df.rename(columns={…})
rename columns; df.columns.str.lower() normalizes them all at once
s.str.contains("x")
vectorized string methods; pass case=False, na=False so missing values don’t break the mask
s.str.extract(r"(\d+)")
regex groups into columns; str.split(",", expand=True) splits into columns
pd.to_numeric(s, errors=)
errors="coerce" turns unparseable strings into NaN instead of raising
s.map({"a": 1})
look each value up in a dict; values without a match become NaN
pd.cut · pd.qcut
bin by fixed edges, e.g. ages · into equal-count quantiles: pd.qcut(s, 4)
df.explode("tags")
one row per element of a list column
astype("category")
10× memory win on low-cardinality strings

Groupby & joins

groupby("k")["v"].mean()
split-apply-combine, the core move
.agg(["mean","count"])
multiple stats at once; dict for per-column stats
.agg(total=("v", "sum"))
named aggregation: flat output columns, one name=(column, func) pair each
.transform("mean")
group stat broadcast back to the original rows — e.g. subtract each group’s mean in one line
as_index=False
with groupby: keep group keys as columns, so no reset_index() afterwards
groupby(dropna=False)
by default rows whose key is NaN silently vanish from the result
pd.merge(a, b, on="k")
SQL joins with how="left" or "inner"; indicator=True debugs match rates
validate="m:1"
make a merge raise when the right key is not unique — catches silent row explosions
pd.concat([a, b])
stack frames; ignore_index=True renumbers, axis=1 goes side by side

Reshape & time series

pivot_table / melt
long→wide with aggregation / wide→long (tidy)
pd.to_datetime + .dt
parse once, then .dt.date/.dt.dayofweek; resample by .resample("W")
resample("W", on="ts")
time buckets from a datetime column; month-end is "ME" — pandas 3 rejects "M"
s.rolling(7).mean()
7-row moving average; rolling("7D") uses a time window on a DatetimeIndex
s.shift(1) · s.diff()
previous row’s value · change from it; pct_change() gives growth
set_index · reset_index
move a column into the index and back; many bugs are a forgotten reset_index()

Speed & memory

df.apply(f, axis=1)
a Python loop in disguise — slow; vectorize with column math or np.where
memory_usage(deep=True)
true memory per column, counting string contents
df.to_numpy()
the underlying array for NumPy or scikit-learn; prefer it to .values

NumPy core

Vectorize, never loop
a * b + c on whole arrays — 100× over Python loops
Broadcasting
shapes align right-to-left; (n,1) × (1,m) → (n,m)
np.where(cond, x, y)
vectorized if/else; chain with np.select
argmax / argsort
indices, not values — top-k = argsort()[-k:][::-1]
np.argpartition(a, -k)
top-k in O(n) without a full sort; the k results come back unordered
@ / np.dot
matrix multiply — cosine sim = a@b / (norm(a)*norm(b))
axis rule
axis=0 collapses rows (per-column), axis=1 collapses columns (per-row)
a.reshape(-1, 1)
-1 infers that dimension — here, a 1-D array becomes a column
np.random.default_rng(0)
the modern seeded generator: rng.normal(size=5), rng.integers(0, 10, 5)
Views vs copies
a slice is a view, so editing it changes the original; .copy() breaks the link
np.nanmean(a)
statistics that skip NaN; plain mean returns NaN if any value is NaN
np.linalg.norm(a, axis=1)
row-wise vector lengths; keepdims=True keeps a column you can divide by
np.unique(a)
sorted distinct values; return_counts=True adds how often each appears
a.astype(np.float32)
half the memory of float64; integer types overflow silently (int8 wraps past 127)

Recipes

Aggregate per group and keep the top 3 in each

import pandas as pd

df = pd.read_csv("sales.csv", parse_dates=["date"])
summary = (
    df.groupby(["region", "product"], as_index=False)
      .agg(revenue=("amount", "sum"), orders=("order_id", "nunique"))
      .sort_values(["region", "revenue"], ascending=[True, False])
      .groupby("region")
      .head(3)
)

Named aggregation gives flat column names, and groupby(...).head(3) after sorting keeps the first three rows per group: the pandas version of SQL’s top-N per group.

Join two tables and check the match rate

merged = orders.merge(customers, on="customer_id", how="left",
                      indicator=True, validate="m:1")
print(merged["_merge"].value_counts())
unmatched = merged[merged["_merge"] == "left_only"]

validate="m:1" raises if a customer appears twice, which would duplicate orders, and the _merge column counts the orders that found no customer.

Weekly totals and a 4-week moving average

weekly = (
    df.set_index("date")["amount"]
      .resample("W")
      .sum()
)
trend = weekly.rolling(4).mean()

resample needs a datetime index (or on="date"), and the first three values of the moving average are NaN because the window is not yet full.

Cosine similarity of one vector against many

import numpy as np

docs = np.random.default_rng(0).normal(size=(1000, 384))
query = docs[0]
docs_n = docs / np.linalg.norm(docs, axis=1, keepdims=True)
scores = docs_n @ (query / np.linalg.norm(query))
top5 = np.argsort(scores)[-5:][::-1]

Normalizing the rows once turns every similarity into a single matrix-vector product; keepdims=True keeps the norms as a column so broadcasting divides each row.

Shrink a DataFrame’s memory

before = df.memory_usage(deep=True).sum()
for col in df.columns:
    if pd.api.types.is_string_dtype(df[col]) and df[col].nunique() < 0.5 * len(df):
        df[col] = df[col].astype("category")
df = df.astype({c: "float32" for c in df.select_dtypes("float64").columns})
print(before, df.memory_usage(deep=True).sum())

Repetitive strings become categories and floats drop to 32 bits; is_string_dtype works on pandas 2 and 3, where string columns have different dtypes.

Questions about Pandas & NumPy Cheat Sheet

What is the difference between loc and iloc in pandas?

loc selects by label: index values and column names, plus boolean masks, and its slices include the end label, so df.loc[2:5] can return four rows. iloc selects by integer position, like a Python list, so df.iloc[2:5] returns the third to fifth rows and excludes position 5. The difference only matters when the index is not a plain 0 to n-1 range, which is exactly when bugs appear, for example after filtering or sorting. Use loc whenever you mean a named row or column, and iloc for first, last or nth.

How do I fix SettingWithCopyWarning in pandas?

The warning appears when you assign through chained indexing, such as df[df.status == "late"]["fee"] = 10, because the first step may return a copy and the assignment would then be lost. The fix is a single indexing call that selects rows and column together: df.loc[df.status == "late", "fee"] = 10. When you really want a separate DataFrame, call .copy() when you create it. In pandas 3, Copy-on-Write is the default, the old warning is gone, and chained assignment never modifies the original, so code that relied on it must switch to loc.

Why is my pandas code slow?

Usually because Python is running once per row. apply with axis=1, iterrows and explicit loops all execute Python code for every row, while column arithmetic, comparisons, np.where and the built-in string and datetime accessors run in compiled code over the whole column at once, often a hundred times faster. Next, check dtypes: numbers stored as strings and object columns are slow and large, so convert them and use category for repetitive text. For big inputs, read only the columns you need and prefer Parquet over CSV.

What is broadcasting in NumPy?

Broadcasting is how NumPy combines arrays of different shapes without copying data. Shapes are compared from the last dimension backwards, and two dimensions are compatible when they are equal or one of them is 1; a dimension of 1 is stretched to match. So an array of shape (1000, 3) minus one of shape (3,) subtracts the same row from all 1000 rows, and shapes (n, 1) and (1, m) combine into (n, m). When shapes do not line up, add an axis with reshape or keepdims rather than writing a loop.