Predictive Modeling and Machine Learning Fundamentals Questions
Applied predictive-modeling cases worked end to end on tabular business data. Covers choosing a model family for a business target (churn, lifetime value, lead scoring, conversion, support-ticket escalation), setting a baseline, defining the label, observation window and prediction horizon, catching leakage, sanity-checking and interpreting results for non-technical stakeholders (feature importance and its traps with correlated features, explaining a single flagged customer), handling censored outcomes with survival methods such as Kaplan-Meier and Cox models, turning model scores into thresholds and business actions with cost-sensitive reasoning, targeting for incremental effect with uplift modeling, and clustering and unsupervised methods in practice: k-means and its initialization (k-means++), empty clusters, scaling, distance choices and mini-batch variants, choosing and validating the number of clusters without labels (elbow, silhouette, stability), density-based, hierarchical and mixture-model alternatives, mixed numeric and categorical features, high-dimensional preparation and visualization, and customer segmentation carried through to stakeholder communication, testing whether segment-based targeting works, keeping segments stable over refits, and downstream use. Excludes regression and classification algorithm internals, metric and validation theory, feature engineering mechanics, and large-scale or production ML systems (including distributed or streaming clustering), which are covered elsewhere.
Your segmentation was built on last year's data. Each time you refit it, the clusters reshuffle, so 'segment 2' last quarter looks like 'segment 4' now, and dashboards and playbooks break. How do you keep segment identities usable across refits?
Sample Answer
The idea in plain words
K-means (and most clustering) numbers its groups arbitrarily. The numbers 0, 1, 2, 3 come from random starting points, so a refit can describe the same four groups with the numbers shuffled. The fix is to stop treating the algorithm's number as the identity of a segment. The segment's identity is a permanent ID that you assign once and then carry across refits by matching new clusters to the old ones.
Three options, in the order I would reach for them
- Freeze the model and only assign. Keep last year's scaler and centroids, assign each new customer to the nearest centroid. IDs cannot move because nothing is refit. This is the right default for dashboards and playbooks. Refit on a schedule (quarterly or yearly) or when a drift check fails (for instance, customers now sit on average much farther from their assigned centre than they did when the model was fitted; the example below prints that number), not on every run.
- Refit, then match to the reference. Compute the distance between every old centroid and every new centroid, then find the one-to-one pairing with the smallest total distance (the Hungarian algorithm, or optimal assignment, implemented as
scipy.optimize.linear_sum_assignment). Give each new cluster the permanent ID of its partner. When the same customers appear in both fits, you can instead pair clusters by the overlap counts in a contingency matrix (rows old IDs, columns new IDs). - Warm-start the refit. Pass last quarter's centroids as the starting points (
init=in scikit-learn'sKMeans, withn_init=1). Labels then carry over by construction, because each centroid keeps drifting from where it was.
Worked example
Four groups on two features. Quarter 1 is fit once and its four centroids become permanent segments 0 to 3. In quarter 2 the groups drift a little (centres move by Gaussian noise with standard deviation 0.5), and a fresh refit with a different random seed is run. Then a fifth group appears and the model is refit with five clusters.
import numpy as np
from scipy.optimize import linear_sum_assignment
from scipy.spatial.distance import cdist
from sklearn.cluster import KMeans
rng = np.random.default_rng(0)
K = 4
base = np.array([[0, 0], [6, 0], [0, 6], [6, 6]], float)
def sample(shift, n=300):
"""One quarter of customers; centers drift by `shift`."""
return np.vstack([c + shift[i] + rng.normal(0, 1.0, (n, 2)) for i, c in enumerate(base)])
# Quarter 1: fit once, give each centroid a permanent business ID by inspection
X1 = sample(np.zeros((K, 2)))
km1 = KMeans(K, n_init=10, random_state=0).fit(X1)
# order the Q1 centroids so that row i is segment i (done once, by a person, from the profile)
order = [int(np.argmin(np.linalg.norm(km1.cluster_centers_ - b, axis=1))) for b in base]
ref_centroids = km1.cluster_centers_[order] # row i = permanent segment i
# Quarter 2: the same four groups, drifted a little. Refit from scratch with a different seed
drift = rng.normal(0, 0.5, (K, 2))
X2 = sample(drift)
km2 = KMeans(K, n_init=1, random_state=5).fit(X2)
print("raw new labels for the 4 groups (arbitrary ids):",
[int(np.bincount(km2.labels_[i*300:(i+1)*300]).argmax()) for i in range(K)])
# Fix 1: match new centroids to reference centroids with the Hungarian algorithm (minimum total distance)
cost = cdist(ref_centroids, km2.cluster_centers_)
rows, cols = linear_sum_assignment(cost)
new_to_ref = {int(c): int(r) for r, c in zip(rows, cols)}
stable = np.array([new_to_ref[l] for l in km2.labels_])
print("new id -> permanent id:", new_to_ref)
print("stable ids for the 4 groups:",
[int(np.bincount(stable[i*300:(i+1)*300]).argmax()) for i in range(K)])
print("matched centroid distances:", np.round(cost[rows, cols], 2), " all under 3.0:", bool((cost[rows, cols] < 3.0).all()))
# Fix 2: warm start the refit from last quarter's centroids so ids carry over by construction
km2w = KMeans(K, init=ref_centroids, n_init=1).fit(X2)
print("warm-start ids for the 4 groups:",
[int(np.bincount(km2w.labels_[i*300:(i+1)*300]).argmax()) for i in range(K)])
print("same partition as matched refit:", bool((km2w.labels_ == stable).all()))
# A real change: a 5th group appears. Matching must leave it unmatched, not steal an old id
X3 = np.vstack([X2, np.array([12.0, 12.0]) + rng.normal(0, 1.0, (300, 2))])
km3 = KMeans(5, n_init=10, random_state=0).fit(X3)
cost3 = cdist(ref_centroids, km3.cluster_centers_)
r3, c3 = linear_sum_assignment(cost3)
pairs = {int(c): int(r) for r, c in zip(r3, c3)}
print("5-cluster refit, new id -> permanent id:", pairs)
extra = [j for j in range(5) if j not in pairs]
print("new cluster with no old partner:", extra, "-> issue a new permanent id (4)")
Output:
raw new labels for the 4 groups (arbitrary ids): [0, 3, 2, 1]
new id -> permanent id: {0: 0, 3: 1, 2: 2, 1: 3}
stable ids for the 4 groups: [0, 1, 2, 3]
matched centroid distances: [1. 0.33 0.89 0.51] all under 3.0: True
warm-start ids for the 4 groups: [0, 1, 2, 3]
same partition as matched refit: True
5-cluster refit, new id -> permanent id: {4: 0, 2: 1, 0: 2, 3: 3}
new cluster with no old partner: [1] -> issue a new permanent id (4)
Reading it: the raw refit labels the four groups [0, 3, 2, 1], so segments 1 and 3 swapped and a dashboard keyed on the number would silently show the wrong group. Matching gives new 3 -> permanent 1 and new 1 -> permanent 3, and the groups read [0, 1, 2, 3] again. The matched distances are all small (largest 1.00), under an example cut-off of 3.0. The warm-started refit gives the same partition as the matched refit (True), which is a useful agreement check but not a general guarantee. In the five-cluster refit the four old IDs are paired with the four nearest new clusters and new cluster 1 has no old partner, so it receives a new permanent ID (4).
Overlap counts, a drift check and a lineage table in miniature
A contingency matrix is a table of counts: rows are the old segment ids, columns the new ones, and each cell counts customers who sit in that old and that new segment. Where the same customers are in both fits, the biggest cell in each row names the partner. A drift check measures whether today's customers still fit the frozen segments, for example the average distance from each customer to its nearest centre, compared with the same number in the quarter the segments were fitted. A lineage table has one row per segment id per release recording where the id came from. The code builds all three on small made-up data; the second part uses two centres and 500 customers each, and the drifted quarter shifts every centre by 1.5 in both features.
import numpy as np
import pandas as pd
from scipy.optimize import linear_sum_assignment
# the same 12 customers in two fits; the numbers are arbitrary labels, so the groups carry different numbers in each fit
old = np.array([0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2, 2])
new = np.array([2, 2, 2, 1, 0, 0, 0, 1, 1, 1, 1, 3])
cm = pd.crosstab(pd.Series(old, name="old id"), pd.Series(new, name="new id"))
print(cm.to_string())
rows, cols = linear_sum_assignment(-cm.values) # maximise the total overlap
print("pairing (new id -> permanent id):", {int(cm.columns[c]): int(cm.index[r]) for r, c in zip(rows, cols)})
print("unpaired new ids:", [int(c) for c in cm.columns if c not in cm.columns[cols]])
# drift check: how far customers sit from their assigned centre now, compared with when the segments were fitted
rng = np.random.default_rng(0)
centres = np.array([[0.0, 0.0], [6.0, 0.0]])
def mean_dist(X):
return np.linalg.norm(X[:, None] - centres[None], axis=2).min(axis=1).mean()
q1 = np.vstack([c + rng.normal(0, 1.0, (500, 2)) for c in centres])
q_same = np.vstack([c + rng.normal(0, 1.0, (500, 2)) for c in centres])
q_moved = np.vstack([c + [1.5, 1.5] + rng.normal(0, 1.0, (500, 2)) for c in centres])
base = mean_dist(q1)
print(f"mean distance to the assigned centre: fit quarter {base:.2f}, next quarter {mean_dist(q_same):.2f}, drifted quarter {mean_dist(q_moved):.2f}")
# a lineage table: one row per permanent id per release
lineage = pd.DataFrame([
dict(release="2026Q3", permanent_id=0, parent_id="none", reason="carried over"),
dict(release="2026Q3", permanent_id=4, parent_id="2", reason="split from 2, new group of larger-basket buyers"),
])
print(lineage.to_string(index=False))
Output:
new id 0 1 2 3
old id
0 0 1 3 0
1 3 0 0 0
2 0 4 0 1
pairing (new id -> permanent id): {2: 0, 0: 1, 1: 2}
unpaired new ids: [3]
mean distance to the assigned centre: fit quarter 1.26, next quarter 1.26, drifted quarter 2.39
release permanent_id parent_id reason
2026Q3 0 none carried over
2026Q3 4 2 split from 2, new group of larger-basket buyers
Reading it: old segment 0 has 3 customers in new id 2, so new 2 becomes permanent 0; old 1 has 3 in new 0, so new 0 becomes permanent 1; old 2 has 4 in new 1, so new 1 becomes permanent 2. New id 3 holds a single customer from old 2 and is left unpaired, which is the signal to issue a new id or treat it as noise at that size. The optimal pairing maximises the total overlap (4 + 3 + 3 = 10 of 12 customers), the same one-to-one logic as matching centroids. For the drift check, the average distance to the assigned centre is 1.26 in the fitting quarter and 1.26 in a next quarter drawn the same way, but 2.39 once the groups have moved; a rule such as "refit when the number rises 20% above the fitting quarter's" (the 20% is a choice, set from how much movement you can tolerate) would fire only on the third. The lineage rows show the minimum to record: release, id, parent, and the reason.
Rules that make it production-safe
- A distance cut-off. A matched pair farther apart than the cut-off is not the same segment even though the algorithm paired it (this matters when a group vanishes and another appears while k stays the same, because one-to-one matching will always pair everything). Choose the cut-off from the spread of a segment, for example a multiple of its typical member distance.
- Retire, split and merge explicitly. If one old segment now has two close new centroids, record it as a split (keep the old ID on the bigger part, issue a new ID for the other). If two old IDs collapse into one, record it as a merge. Keep a small lineage table (new ID, parent ID, date, reason).
- Name segments by a rule, not an index. Store a one-line definition and a profile per ID; write playbooks against that. After each refit, check that the matched segment still has the old profile (size, key feature means).
- Version the model. Every customer row carries the model version, so last quarter's report can be reproduced.
- Allow some customers to change segment and report the share that moved; a very high churn of membership means the segments are not stable enough to act on.
Trade-offs
Freezing is simplest and keeps history comparable, but segments go stale as behaviour drifts. Refit with matching tracks the population but needs the lineage and cut-off logic, and the segment's meaning can drift slowly while its ID stays the same. Warm-start is neat but can lock a bad solution in place when groups truly change. My default: freeze and assign daily, refit with matching on a quarterly schedule or on a drift alarm, and review the lineage by a human before the new version is released.
How does your approach to building a customer segmentation change if it is a one-time report on a static snapshot versus something that must assign new customers to segments every day?
Sample Answer
Direct answer
The main change is that a daily job cannot refit the segmentation every time. The model becomes a frozen artefact (the preprocessing steps plus the cluster centres) and each new customer is assigned to the nearest centre. A one-time report can be fit, read, and thrown away; a daily system needs stable definitions, a deterministic assignment step, monitoring for drift, and a scheduled refit.
Static report versus daily assignment
| Concern | One-time snapshot | Daily assignment |
|---|---|---|
| Fitting | Fit once, try several settings freely | Fit offline on a schedule, then freeze |
| Preprocessing | Anything goes | Same scaler, same transforms, same missing-value rules saved with the model |
| Feature choice | Whatever is in the snapshot | Only features available for a new customer on day one (a new customer has little history) |
| Segment identity | Numbers only matter inside the report | Permanent IDs and written definitions that stay put |
| Quality | Check once | Monitor share per segment, distance to the centre, and unusual cases |
| Failure | A wrong slide | A wrong action sent to many customers every day |
Mechanism of the daily step
Scale the new customer's features with the training mean and standard deviation (the scaler is saved, never refit), compute the distance to each centroid, and take the nearest one. The minimum distance is also a free health signal: an unusually large value means the customer fits no segment well.
Worked example
A snapshot of 1,000 customers in three segments (orders per year, average basket, days since last order). The scaler and k-means are fit once. Then two daily batches of 100: day A from the same population, day B in which 60 customers are drawn at random from the day A batch and 40 are a new kind of customer (frequent buyers with tiny baskets). The limit is the 99th percentile of the training distances.
import numpy as np
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
rng = np.random.default_rng(0)
# Snapshot: 3 segments over 3 features (orders per year, average basket, days since last order)
centers = np.array([[2, 30, 200], [12, 45, 20], [30, 25, 5]], float)
sizes = [500, 300, 200]
snap = np.vstack([c + rng.normal(0, [1.5, 5, 25 if i == 0 else 8], (n, 3))
for i, (c, n) in enumerate(zip(centers, sizes))])
# One-time report: fit, read the profile, done.
scaler = StandardScaler().fit(snap) # frozen: the scaling is part of the model
km = KMeans(3, n_init=10, random_state=0).fit(scaler.transform(snap))
train_dist = np.linalg.norm(scaler.transform(snap) - km.cluster_centers_[km.labels_], axis=1)
train_share = np.bincount(km.labels_, minlength=3) / len(snap)
limit = np.quantile(train_dist, 0.99) # "unusually far from every centroid" cut-off
def assign_daily(batch):
"""Daily job: reuse the frozen scaler and centroids. Nothing is refit."""
Z = scaler.transform(batch)
d = np.linalg.norm(Z[:, None, :] - km.cluster_centers_[None], axis=2) # (n, 3)
return d.argmin(axis=1), d.min(axis=1)
def report(name, batch):
lab, dmin = assign_daily(batch)
share = np.bincount(lab, minlength=3) / len(batch)
print(f"{name}: segment shares {np.round(share, 2)} "
f"mean distance {dmin.mean():.2f} share beyond 99th-pct limit {(dmin > limit).mean():.3f}")
print("training: customers per k-means label", np.bincount(km.labels_))
print("training: segment shares", np.round(train_share, 2),
" mean distance", round(float(train_dist.mean()), 2), " limit", round(float(limit), 2))
# Day A: same population as training. Day B: a new kind of customer (frequent buyers with tiny baskets) appears
dayA = np.vstack([c + rng.normal(0, [1.5, 5, 25 if i == 0 else 8], (n, 3))
for i, (c, n) in enumerate(zip(centers, [50, 30, 20]))])
newgroup = np.array([25, 8, 10]) + rng.normal(0, [1.5, 2, 8], (40, 3))
old60 = dayA[rng.choice(len(dayA), 60, replace=False)] # 60 random customers from the usual population
dayB = np.vstack([old60, newgroup])
report("day A", dayA)
report("day B", dayB)
print("day B: the 40 new customers land in labels", np.bincount(assign_daily(newgroup)[0], minlength=3),
" the 60 usual customers land in", np.bincount(assign_daily(old60)[0], minlength=3))
Output:
training: customers per k-means label [300 500 200]
training: segment shares [0.3 0.5 0.2] mean distance 0.5 limit 1.35
day A: segment shares [0.3 0.5 0.2] mean distance 0.54 share beyond 99th-pct limit 0.020
day B: segment shares [0.19 0.31 0.5 ] mean distance 1.04 share beyond 99th-pct limit 0.400
day B: the 40 new customers land in labels [ 0 0 40] the 60 usual customers land in [19 31 10]
Reading it: k-means numbers its groups in its own order, not in the order the code built them, so label 0 is the 300-customer group, label 1 the 500-customer group and label 2 the 200-customer group (the first printed line). That is why the training shares read 0.3, 0.5, 0.2 even though the code lists the sizes as 500, 300, 200.
Day A is a fresh draw from the same population, so its shares match training and the mean distance (0.54) is close to 0.50. About 2% of day A lies beyond the limit. The limit is the 99th percentile of the training distances, so by construction about 1% of ordinary customers sit beyond it; a few percent is normal, tens of percent is not.
Day B mixes 60 customers drawn at random from the usual population with 40 customers of a new kind (frequent buyers with tiny baskets). The usual 60 spread over the labels as 19, 31 and 10, which is what the first two shares (0.19, 0.31) show. The new 40 all land in label 2, because it is the nearest centroid, so label 2 jumps from 0.2 to 0.5. None of those 40 belong there: the last printed line shows 10 usual customers and 40 new ones. The shares alone would suggest that segment 2 simply grew; the distance check shows what really happened. The mean distance roughly doubles (1.04 against 0.50) and 40% of the batch is beyond the limit, because the new customers are far from every centroid. Drift here means that the incoming population no longer looks like the one the segments were fit on. That is the case for a refit or a new segment, not for acting on the shares.
What I would put around it
- Save the scaler and centroids together as a versioned artefact; write the model version next to each assignment.
- Track the share per segment and the distance distribution daily, and alert on large moves.
- Refit on a schedule or when the alert fires, and match new clusters to the old IDs so dashboards stay intact.
- Handle customers with almost no history: either a "not enough data yet" label or a first-order model on the few features known at signup.
Flip condition
If the audience is a single decision meeting, spend the effort on interpretation and stability of the segments, not on serving infrastructure. If assignments drive automated actions, spend it on monitoring.
You are asked to predict which support tickets will escalate. How do you define the target, deal with labels that only arrive weeks later, and evaluate the model before anyone relies on it?
Sample Answer
Three decisions, in this order: define exactly what event the model predicts and at what moment, train only on tickets whose outcome has had time to arrive, and judge the model on later tickets with a metric that matches how the predictions will be used. The delayed label is the hazard that is specific to this problem: a ticket that has not escalated yet is not a negative, it is an unknown.
1. Define the target
Write the target as a sentence with four parts: prediction moment, event, horizon, population.
- Prediction moment: when the score is needed. Say "at ticket creation" if routing or staffing decisions happen then, or "after the first agent reply" if the score drives a later review queue. Every feature must be available at that moment, no later.
- Event: what counts as escalation, as the support system actually records it (for example, moved to a senior tier or flagged by a manager). Check that the event is stable: if teams escalate by different rules, or a policy change redefined escalation in the middle of the history, the label means different things in different periods.
- Horizon: "escalates within 45 days of creation". Without a horizon every ticket eventually has some chance of escalating and the label depends on how long you waited. The horizon comes from the decision: if reviewers act within a week, a 7-day horizon may be the useful target even if some tickets escalate later.
- Population: exclude tickets that were escalated by an automatic routing rule at creation, or spam, because predicting those teaches nothing.
Name the decision the score supports (route to a senior queue, reserve reviewer time, alert a manager) because that decides the metric in step 3 and the threshold.
2. Labels that arrive weeks later
A ticket created 5 days ago that has not escalated might still. If you label it 0 you are training on false negatives, concentrated among the newest tickets, so the model learns that recent-looking patterns are safe. This is right-censoring: for a recent ticket you only watched part of the horizon, because today's date cuts the observation window short, so you know the event had not happened yet but not whether it will. Options, in the order I would reach for them:
- Maturity window. Train only on tickets old enough for the full horizon to have passed at the training snapshot (created at least 45 days earlier). This is simple and unbiased for the label, at the cost that the model is always trained on data at least one horizon old. Choose the horizon from the observed delay distribution, so you know how much escalation falls after it.
- Survival or time-to-event model (a model of "how long until the event, with some tickets still waiting"; a Cox model is the best-known one) if you need to use immature tickets or need a risk curve over time. More machinery, and it needs assumptions such as proportional hazards (the effect of each feature multiplies the escalation rate by a constant factor at every point in time).
- Early proxy label (for example, "customer asked for a manager within 48 hours") only if you can show on matured tickets that it tracks the real label, and you re-check that regularly. State it as a proxy, never as the target.
Also decide how the model stays current: because labels lag, retrain on a schedule and monitor input drift (the distribution of incoming features moving away from what the model was trained on, which you can see without waiting for labels) rather than waiting for labels to reveal a problem.
3. Evaluate before anyone relies on it
The bullets are in the order I would do them; if time is short, the time split, the metric from the decision and the leakage audit matter most.
- Split by time, not at random. Train on tickets created before a cutoff, test on tickets created after it whose horizon has fully elapsed. A random split lets near-identical tickets from the same incident land on both sides and flatters the model.
- Pick the metric from the decision. If reviewers can read 10% of tickets, report precision in the top 10% of scores and the lift over the base rate (the overall escalation rate; lift is precision divided by it). If the score is used as a probability to set thresholds or staffing, also check calibration (the probabilities mean what they say: do tickets scored 0.2 escalate about 20% of the time?). AUC (the chance a random escalating ticket scores above a random non-escalating one) summarises ranking only.
- Compare with a simple baseline such as "urgent priority" or "enterprise customer". The model must beat what a rule already gives for free.
- Audit features for leakage. For every feature ask whether it exists at the prediction moment. Fields that are filled in as a ticket moves towards escalation (assignee changes, escalation notes, hand-off counts) make the offline score look excellent and the live one useless.
- Slice by customer tier, product, language and channel, since an average hides groups where it fails.
- Shadow mode before reliance (running the model without letting it act). Run the model silently on live tickets, log scores, wait one horizon, then compare with what happened. Only then let it drive actions, starting with low-cost ones.
A computed example (synthetic data)
10,000 simulated tickets with an urgency level, an enterprise flag and a text score. Escalation delay (creation to escalation) follows a lognormal distribution with a median of 12 days, so a few escalations arrive after the 45-day horizon (a lognormal delay is one whose logarithm is bell-shaped, which gives a long right tail: most escalations come quickly and a few come very late). The training snapshot is day 150, the test set is tickets created after day 150 whose 45 days have elapsed. "Naive" labels mean "escalated by the snapshot". "Horizon-matured" labels use only tickets created at least 45 days before the snapshot. The "horizon rate" column is the true 45-day label, which is only known here because the data is simulated.
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
rng = np.random.default_rng(0)
N, T_TRAIN, TODAY, H = 10000, 150, 240, 45 # tickets, training snapshot day, today, label horizon (days)
created = rng.uniform(0, 220, N) # creation day
prio = rng.integers(0, 3, N) # 0 low, 1 normal, 2 urgent
enterprise = rng.integers(0, 2, N)
text = rng.normal(0, 1, N) # stand-in for a text-derived score
p = 1 / (1 + np.exp(-(-3.0 + 0.9 * prio + 0.8 * enterprise + 0.5 * text)))
will_escalate = rng.random(N) < p # ever escalates
delay = np.exp(rng.normal(np.log(12), 0.8, N)) # days from creation to escalation
esc_day = np.where(will_escalate, created + delay, np.inf)
within_h = will_escalate & (delay <= H) # the target: escalates within H days of creation
X = np.c_[prio, enterprise, text]
print("share of escalations that arrive after day", H, "from creation:", round(float((delay[will_escalate] > H).mean()), 3))
# what a naive label looks like at the training snapshot: "escalated so far"
known_by_T = created <= T_TRAIN
naive = known_by_T & (esc_day <= T_TRAIN)
for lo, hi in [(0, 15), (15, 30), (30, 45), (45, 400)]:
m = known_by_T & ((T_TRAIN - created) >= lo) & ((T_TRAIN - created) < hi)
label = f"{lo}-{hi}" if hi < 400 else f"{lo}+"
print(f"age {label:>5} days at snapshot: naive rate {naive[m].mean():.3f} horizon rate {within_h[m].mean():.3f}")
# time-based evaluation: test = tickets created after the snapshot whose horizon has fully elapsed by TODAY
test = (created > T_TRAIN) & (created <= TODAY - H)
mature = known_by_T & (created <= T_TRAIN - H) # only tickets whose horizon had elapsed at the snapshot
fit = lambda tr, y, F: LogisticRegression().fit(F[tr], y[tr])
m_naive = fit(known_by_T, naive, X)
m_ok = fit(mature, within_h, X)
cap = int(0.10 * test.sum()) # reviewers can read 10% of tickets
print("test tickets:", int(test.sum()), " actual rate:", round(float(within_h[test].mean()), 3))
for name, m in [("naive labels", m_naive), ("horizon-matured labels", m_ok)]:
s = m.predict_proba(X[test])[:, 1]
top = np.argsort(-s)[:cap]
print(f"{name:22s} mean predicted {s.mean():.3f} AUC {roc_auc_score(within_h[test], s):.3f} precision in top 10% {within_h[test][top].mean():.3f} lift {within_h[test][top].mean() / within_h[test].mean():.2f}x")
# a leaky feature: number of hand-offs between agents, only recorded as the ticket moves toward escalation
handoffs = rng.poisson(0.3 + 2.0 * within_h)
XL = np.c_[X, handoffs]
m_leak = fit(mature, within_h, XL)
print("with a hand-off count that does not exist at prediction time: AUC", round(roc_auc_score(within_h[test], m_leak.predict_proba(XL[test])[:, 1]), 3))
# the same comparison at three different training snapshots: the earlier the snapshot, the more immature tickets there are
print("\nsnapshot immature share actual rate naive: mean predicted / AUC matured: mean predicted / AUC")
for T in (50, 90, 150):
kb = created <= T
nv = kb & (esc_day <= T)
mat = kb & (created <= T - H)
te = (created > T) & (created <= TODAY - H)
res = []
for tr, lab in [(kb, nv), (mat, within_h)]:
sc = fit(tr, lab, X).predict_proba(X[te])[:, 1]
res.append((sc.mean(), roc_auc_score(within_h[te], sc)))
print(f"day {T:3d} {1 - mat.sum() / kb.sum():14.2f} {within_h[te].mean():11.3f} "
f"{res[0][0]:19.3f} / {res[0][1]:.3f} {res[1][0]:19.3f} / {res[1][1]:.3f}")
# the three-snapshot comparison again on 20 other simulated sets of tickets (same generator, new seeds):
# actual rate minus the mean predicted probability, and the AUC difference naive minus matured
def gaps(seed):
r = np.random.default_rng(seed)
cr = r.uniform(0, 220, N); pr = r.integers(0, 3, N); en = r.integers(0, 2, N); tx = r.normal(0, 1, N)
wl = r.random(N) < 1 / (1 + np.exp(-(-3.0 + 0.9 * pr + 0.8 * en + 0.5 * tx)))
dl = np.exp(r.normal(np.log(12), 0.8, N))
ed = np.where(wl, cr + dl, np.inf)
wh = wl & (dl <= H)
F = np.c_[pr, en, tx]
out = []
for T in (50, 90, 150):
kb = cr <= T
te = (cr > T) & (cr <= TODAY - H)
sc = [fit(tr, lab, F).predict_proba(F[te])[:, 1] for tr, lab in [(kb, kb & (ed <= T)), (kb & (cr <= T - H), wh)]]
out.append([wh[te].mean() - sc[0].mean(), wh[te].mean() - sc[1].mean(),
roc_auc_score(wh[te], sc[0]) - roc_auc_score(wh[te], sc[1])])
return out
G = np.array([gaps(sd) for sd in range(1, 21)]) # shape (20 sets, 3 snapshots, 3 numbers)
print("\n20 other simulated sets: actual rate minus mean predicted (mean, then min to max), and AUC naive minus matured (min to max)")
for t, T in enumerate((50, 90, 150)):
g = G[:, t]
print(f"day {T:3d} naive {g[:, 0].mean():+.3f} ({g[:, 0].min():+.3f} to {g[:, 0].max():+.3f}) "
f"matured {g[:, 1].mean():+.3f} ({g[:, 1].min():+.3f} to {g[:, 1].max():+.3f}) AUC diff {g[:, 2].min():+.3f} to {g[:, 2].max():+.3f}")
Output:
share of escalations that arrive after day 45 from creation: 0.047
age 0-15 days at snapshot: naive rate 0.068 horizon rate 0.206
age 15-30 days at snapshot: naive rate 0.131 horizon rate 0.166
age 30-45 days at snapshot: naive rate 0.184 horizon rate 0.190
age 45+ days at snapshot: naive rate 0.189 horizon rate 0.181
test tickets: 2077 actual rate: 0.192
naive labels mean predicted 0.170 AUC 0.721 precision in top 10% 0.493 lift 2.57x
horizon-matured labels mean predicted 0.179 AUC 0.720 precision in top 10% 0.488 lift 2.54x
with a hand-off count that does not exist at prediction time: AUC 0.907
snapshot immature share actual rate naive: mean predicted / AUC matured: mean predicted / AUC
day 50 0.90 0.188 0.129 / 0.723 0.172 / 0.717
day 90 0.49 0.190 0.153 / 0.723 0.176 / 0.722
day 150 0.30 0.192 0.170 / 0.721 0.179 / 0.720
20 other simulated sets: actual rate minus mean predicted (mean, then min to max), and AUC naive minus matured (min to max)
day 50 naive +0.052 (+0.029 to +0.067) matured -0.000 (-0.055 to +0.052) AUC diff -0.003 to +0.032
day 90 naive +0.027 (+0.006 to +0.039) matured +0.001 (-0.028 to +0.013) AUC diff -0.002 to +0.001
day 150 naive +0.013 (-0.007 to +0.029) matured +0.001 (-0.024 to +0.014) AUC diff -0.002 to +0.002
Reading the run:
- The naive label rate is 0.068 for tickets younger than 15 days at the snapshot against a true horizon rate of 0.206. The naive and true rates agree closely only for tickets at least 30 days old, by which time most escalations have arrived. That gap is the censoring problem made visible.
- Ranking survives, probabilities do not. At the day-150 snapshot the naive model and the matured model rank test tickets almost identically (AUC 0.721 against 0.720, both far above 0.5 but modest), because the missed escalations are spread fairly evenly over the feature values, so the order of tickets barely changes. What the censoring damages is the level of the probabilities: the naive model's mean predicted probability is 0.170 and the matured model's is 0.179, against an actual rate of 0.192 in the test period.
- The last table shows that the level error grows with the share of immature training tickets (those whose horizon has not yet elapsed at the snapshot). At day 50, 90% of the training tickets are immature and the naive model predicts 0.129 on average against an actual 0.188, an under-prediction of about a third; at day 90 it is 0.153 (49% immature), at day 150 0.170 (30% immature). The matured model is much closer to the actual rate (0.172, 0.176, 0.179 against 0.188, 0.190, 0.192), and its AUC is almost unchanged (0.717 to 0.722). Those are the numbers for one simulated set. The last block of the output repeats the comparison on 20 other sets from the same generator: the naive model under-predicts on average by 0.052 at day 50, 0.027 at day 90 and 0.013 at day 150 (at day 150 a few sets come out slightly over), while the matured model has no systematic bias (average gap about 0.000 to 0.001) but is noisy when it has little to learn from. At day 50 it is trained only on tickets created in the first five days, and its level lands anywhere from 0.055 above to 0.052 below the actual rate; at days 90 and 150 it stays within about 0.03. The AUC of the two schemes differs by at most 0.002 at days 90 and 150, and by up to 0.03 at day 50, where either can win. So for this toy the practical answer is: if the score is only used to rank a review queue, the censoring error is mostly harmless; if the probability is used to set thresholds, staffing or an expected-cost calculation, it is not, and the error is worst exactly when the model is young and most tickets are recent. In real data, immature tickets can also differ systematically in their features (new products, new customer types), which can distort the ranking too; this synthetic data does not include that, so the check on your own data is to compare the two label schemes as the table does.
- Precision in the top 10% of scores is 0.488 for the matured model, a lift of 2.54x over the actual rate. That is the number to compare with the reviewers' capacity and with a rule baseline.
- The leaky column is
handoffs = rng.poisson(0.3 + 2.0 * within_h): a count that averages 0.3 for tickets that will not escalate and 2.3 for those that will, which is how a field written during escalation looks. Adding a hand-off count that is only recorded as the ticket moves towards escalation lifts AUC to 0.907. Nothing about the model improved; the feature is a consequence of the outcome. That is why the leakage audit comes before any metric is trusted.
Common mistakes
- Treating "not escalated yet" as a negative.
- Choosing a horizon by convenience instead of from the delay distribution and the decision.
- Random-split evaluation, and features taken from the ticket's full history.
- Reporting AUC alone when the use is a short review queue.
- Going live without a shadow period.
What would change the plan: if the horizon relevant to the decision is a few days, the maturity window shrinks and a simple retrain cadence is enough. If escalations are rare (under about 1%), report precision and recall at fixed capacity with confidence intervals, because small counts make every metric noisy.
You need to predict customer churn for a subscription product. You can choose between a regularized logistic regression, a tree ensemble, and a small neural network. How do you decide which one to ship, and what would you build first as a baseline?
Sample Answer
Direct answer
Choose by the decision the score feeds, not by algorithm preference. Build a ladder of baselines first (a no-model base rate, a one-line business rule, then a regularized logistic regression), then challenge the best of them with a gradient-boosted tree ensemble. A small neural network is the last thing to try on a table of customer features. Ship the simplest model that beats the baseline on a time-ordered holdout by more than the run-to-run noise, on the metric that matches the action (for example, how many real churners are in the slice of users you can afford to contact), and whose probabilities are trustworthy enough to put a price on (calibrated: among customers scored near 20% risk, about 20% really churn).
Terms used below
- Churn: a subscriber cancels or fails to renew within a fixed window after the date the score is computed.
- Regularized logistic regression: a linear model of the log-odds of churn. The L2 penalty shrinks coefficients so that correlated features do not produce huge offsetting weights.
- Tree ensemble (gradient boosting): many shallow decision trees built one after another, each fitted to the mistakes of the ones before. It finds thresholds and interactions without hand-made columns.
- Small neural network: layers of weighted sums with nonlinearities between them. Flexible, but on a table of a few dozen columns it needs scaling, tuning and more data to match boosted trees.
- AUC: the chance that a random churner gets a higher score than a random non-churner. PR-AUC (average precision) summarizes precision against recall and is far more sensitive to the top of the ranking when churn is rare. Brier score: mean squared error of the predicted probabilities, lower is better, and it rewards honest probabilities.
- Calibration: whether scores behave like probabilities. A calibration check groups customers by score and compares the average predicted risk with the share who actually churned. A calibration step is a small correction fitted on held-out data that maps raw scores to observed churn rates (two common ones are Platt scaling, a logistic curve fitted to the scores, and isotonic regression, a monotone step function).
- Reason codes and attribution tooling: a reason code is a short, plain reason for one customer's score (for example "logins fell by 7"). A linear model gives them directly from coefficient times feature value; for a tree ensemble you need attribution methods (for example SHAP values) that split one score into per-feature contributions.
- Time-ordered folds: several train/test splits that all train on earlier months and test on the next block, sliding forward, so each test block is in the future of its training data.
The baseline ladder (what to build first)
- Base rate. Everyone gets the training churn rate. AUC is 0.5 by construction. This sets the floor for every metric, including Brier.
- One-line rule, such as "fewer logins in the last 30 days means riskier". Stakeholders already believe it, so any model has to beat it to justify its cost.
- Regularized logistic regression on a handful of clean features. This is the real baseline: it trains in seconds, its coefficients can be read aloud, and its probabilities are often reasonably calibrated before any extra step.
Only then add the challengers: a tree ensemble, and a neural network if there is a reason (raw text, event sequences, or embeddings that a table does not capture).
How to decide among the three
| Criterion | Regularized logistic | Tree ensemble | Small neural network |
|---|---|---|---|
| Accuracy on tabular churn data | Good when effects are roughly additive on the log-odds scale. Interactions need hand-built columns | Often the strongest, since thresholds and interactions are found automatically | At best comparable, with more tuning |
| Explaining a score to a retention team | Coefficients map to plain sentences | Needs extra attribution tooling (see Terms) | Hardest |
| Probability quality | Often decent as trained | Often needs a calibration step, so check it rather than assume it | Often needs one |
| Cost to maintain and monitor | Lowest | Moderate | Highest |
The deciding rule: take the challenger only if its lead on the action metric survives repeated time-ordered folds (train on earlier months, test on the next one, slide forward) and the team can monitor and explain it. A lead that appears on one split and vanishes on the next is noise.
Worked example
The table below is simulated so that every number can be reproduced. One simulated fact matters for reading the result: new subscribers who also log in rarely are much more likely to churn, an interaction between two columns that a plain logistic regression cannot represent unless someone builds that column. The first 14,000 rows stand in for the earlier period and the last 6,000 for the later period used as the test. After the main table the program also adds the interaction as an explicit logistic column, checks calibration in the top slice, and slides a time-ordered fold forward four times to measure how much the boosted-tree lead moves.
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline
from sklearn.metrics import roc_auc_score, average_precision_score, brier_score_loss
rng = np.random.default_rng(0)
n = 20000
tenure = rng.integers(1, 48, n) # months subscribed
logins = rng.poisson(8, n).astype(float) # logins in last 30 days
logins_prev = rng.poisson(8, n).astype(float) # logins in the 30 days before
tickets = rng.poisson(0.5, n).astype(float) # support tickets last 90 days
late_pay = rng.binomial(1, 0.1, n).astype(float) # a failed payment recently
X = np.column_stack([tenure, logins, logins_prev, tickets, late_pay])
drop = logins_prev - logins # fall in engagement
z = (-3.4 + 0.25 * np.maximum(drop, 0) + 0.5 * tickets + 1.2 * late_pay
- 0.02 * tenure + 2.5 * ((tenure < 6) & (logins < 6))) # new + disengaged interaction
y = rng.binomial(1, 1 / (1 + np.exp(-z)))
split = 14000 # rows are ordered; later rows = test
Xtr, Xte, ytr, yte = X[:split], X[split:], y[:split], y[split:]
print("churn rate train/test:", round(ytr.mean(), 3), round(yte.mean(), 3))
def report(name, p, is_prob=True):
brier = f"{brier_score_loss(yte, p):.4f}" if is_prob else "n/a"
print(f"{name:26s} AUC={roc_auc_score(yte, p):.3f} PR-AUC={average_precision_score(yte, p):.3f} Brier={brier}")
base_rate = np.full(len(yte), ytr.mean())
report("base rate only", base_rate)
rule_score = 1.0 / (1.0 + Xte[:, 1]) # rule: fewer logins = riskier (AUC is rank-only; not a calibrated probability)
report("rule: fewer logins", rule_score, is_prob=False)
lr = make_pipeline(StandardScaler(), LogisticRegression(C=1.0, max_iter=1000)).fit(Xtr, ytr)
report("logistic (L2, C=1)", lr.predict_proba(Xte)[:, 1])
hgb = HistGradientBoostingClassifier(max_depth=3, learning_rate=0.05, max_iter=200, random_state=0).fit(Xtr, ytr)
report("gradient boosted trees", hgb.predict_proba(Xte)[:, 1])
mlp = make_pipeline(StandardScaler(), MLPClassifier((16,), alpha=1e-2, max_iter=300, random_state=0)).fit(Xtr, ytr)
report("small neural net (16)", mlp.predict_proba(Xte)[:, 1])
# --- logistic regression with the interaction added as an explicit column ---
def add_interaction(A):
return np.column_stack([A, (A[:, 0] < 6) & (A[:, 1] < 6)])
lrx = make_pipeline(StandardScaler(), LogisticRegression(C=1.0, max_iter=1000)).fit(add_interaction(Xtr), ytr)
report("logistic + interaction col", lrx.predict_proba(add_interaction(Xte))[:, 1])
# --- calibration: in the 10% of users scored highest, does mean predicted risk match actual churn? ---
print("\ncalibration in the top 10% by score (mean predicted vs actual churn rate):")
for name, p in (("logistic", lr.predict_proba(Xte)[:, 1]), ("gradient boosted trees", hgb.predict_proba(Xte)[:, 1])):
top = np.argsort(-p)[: len(p) // 10]
act = yte[top].mean()
se = (act * (1 - act) / len(top)) ** 0.5 # sampling noise of an observed rate in a slice this size
print(f" {name:24s} predicted={p[top].mean():.3f} actual={act:.3f} (+/- {se:.3f}) (all users: predicted={p.mean():.3f} actual={yte.mean():.3f})")
# --- run-to-run noise: slide a time-ordered fold forward (train on everything earlier, test on the next 4,000 rows) ---
print("\nfold train rows PR-AUC logistic PR-AUC boosted lead of boosted")
leads = []
for start in (4000, 8000, 12000, 16000):
tr, te = slice(0, start), slice(start, start + 4000)
a = make_pipeline(StandardScaler(), LogisticRegression(max_iter=1000)).fit(X[tr], y[tr])
b = HistGradientBoostingClassifier(max_depth=3, learning_rate=0.05, max_iter=200, random_state=0).fit(X[tr], y[tr])
pa = average_precision_score(y[te], a.predict_proba(X[te])[:, 1])
pb = average_precision_score(y[te], b.predict_proba(X[te])[:, 1])
leads.append(pb - pa)
print(f"{start // 4000:4d} {start:10d} {pa:15.3f} {pb:14.3f} {pb - pa:15.3f}")
print(f"lead of boosted trees over the folds: mean {np.mean(leads):.3f}, smallest {min(leads):.3f}, largest {max(leads):.3f}")
churn rate train/test: 0.07 0.064
base rate only AUC=0.500 PR-AUC=0.064 Brier=0.0602
rule: fewer logins AUC=0.663 PR-AUC=0.119 Brier=n/a
logistic (L2, C=1) AUC=0.748 PR-AUC=0.260 Brier=0.0545
gradient boosted trees AUC=0.768 PR-AUC=0.336 Brier=0.0513
small neural net (16) AUC=0.752 PR-AUC=0.275 Brier=0.0537
logistic + interaction col AUC=0.764 PR-AUC=0.340 Brier=0.0510
calibration in the top 10% by score (mean predicted vs actual churn rate):
logistic predicted=0.233 actual=0.237 (+/- 0.017) (all users: predicted=0.070 actual=0.064)
gradient boosted trees predicted=0.252 actual=0.270 (+/- 0.018) (all users: predicted=0.070 actual=0.064)
fold train rows PR-AUC logistic PR-AUC boosted lead of boosted
1 4000 0.215 0.283 0.068
2 8000 0.273 0.313 0.039
3 12000 0.267 0.318 0.051
4 16000 0.264 0.327 0.063
lead of boosted trees over the folds: mean 0.055, smallest 0.039, largest 0.068
How to read it. The no-information model sits at AUC 0.500 with Brier 0.0602, and the logistic regression improves Brier to 0.0545. The tree ensemble leads on this table: AUC 0.768 against 0.748, and PR-AUC 0.336 against 0.260. That PR-AUC gap is the one that matters if the retention program contacts only a small top slice of users. Part of the reason is the built-in interaction, so the size of the gap on real data is an empirical question. The logistic (0.748) and neural network (0.752) AUCs differ by 0.004 on a single seed and a single split, which is not evidence that either is better.
Three further lines of output settle the points the table leaves open.
- Interaction column. Adding the explicit column ((tenure < 6) and (logins < 6)) lifts logistic regression to AUC 0.764, PR-AUC 0.340 and Brier 0.0510, level with the boosted trees (0.768, 0.336, 0.0513). So here it recovers the whole gap, which is expected because the simulation has exactly one interaction and we knew where to look. On real data you rarely know which interaction to add, and finding them automatically is what the trees buy you.
- Calibration. Among the 10% of users scored highest (600 users), logistic predicts 0.233 and 0.237 actually churn; the boosted trees predict 0.252 and 0.270 churn. The two gaps, 0.004 and 0.018, are each about one standard error of an observed rate in a slice this size (the program prints +/- 0.017 and +/- 0.018), so this run does not show that either model is miscalibrated. Drawing the simulation again moves each gap by a few hundredths in either direction, and which model looks better calibrated in the top slice changes with it. The practical point is that a calibration check on a slice of a few hundred users cannot resolve a gap of a couple of points: pool the time-ordered test blocks, or use a larger slice, before concluding that a calibration step is needed. The trees have the better overall Brier score because Brier rewards sharper ranking as well as honest levels. If a gap that survives that larger check matters for pricing, fit a calibration step on held-out data and re-check.
- Run-to-run noise. Across four time-ordered folds the boosted-tree PR-AUC lead over logistic is 0.068, 0.039, 0.051 and 0.063: positive in all four here, but it moves by about 0.03 from fold to fold, and on other draws of the simulation the smallest fold lead can fall to roughly zero. So "more than the run-to-run noise" means a lead that stays clearly positive in every fold, not one that is only large on average. A PR-AUC difference of a few thousandths on one split is well inside that spread, and the 0.004 AUC gap between logistic and the neural network above is the same kind of single-split difference: nothing in this run shows it would repeat.
On this table the call is: ship the tree ensemble if its PR-AUC lead repeats across several time-ordered folds and the team can explain it and monitor it. If the lead shrinks to noise, ship the logistic regression with the interaction added as an explicit column, which, as the output above shows, closed the gap on this table and keeps the explanation simple.
What would flip the choice
- A requirement to give each customer reason codes, or an audit that demands a transparent model, favors logistic regression.
- Very few churners in training (a few hundred) makes flexible models unstable; favor regularized logistic regression.
- Strong interactions or thresholds in the data, and plenty of churners, favor boosted trees.
- Raw sequences, text or images as inputs are the case for a neural network.
Common mistakes
- Choosing on a random split instead of a later period, which flatters every model.
- Reporting overall AUC when the business contacts only the top few percent; report the metric at the contact capacity.
- Skipping calibration when the probabilities feed an expected-value calculation.
- Jumping to a neural network without a baseline, so nobody knows what the complexity bought.
- Comparing models on different feature sets or after tuning on the test period.
Implement k-means from scratch in Python with NumPy. Handle initialization, a maximum number of iterations and a convergence tolerance, and return the centroids and assignments.
Sample Answer
Direct answer
K-means (also called Lloyd's algorithm, the name for this two-step loop) alternates two steps until nothing moves: assign every point to its nearest centroid (a centroid is the mean point of a cluster), then move every centroid to the mean of the points assigned to it. Each step cannot increase the total squared distance from points to their centroids (the inertia), so the loop converges, to a fixed point (a state where one more round changes nothing) that is only a local optimum, one that depends on the starting centroids. The implementation below takes the number of clusters k, an init choice (random points or k-means++), max_iter, a convergence tolerance tol, and returns the centroids, the assignments, the inertia and the iteration count.
The algorithm, step by step
With n points of dimension d stored in X of shape (n,d) and k centroids in C of shape (k,d):
- Initialise
C: either k distinct random data points (random), or k-means++, which picks the first centroid at random and each next one with probability proportional to the squared distance to the nearest centroid already chosen, so starts are spread out. - Assignment step (nearest centroid). Build the (n,k) matrix of squared distances and take
argminalong the centroid axis (for each row, that is, each point, the index of the smallest entry across the centroid columns). The vectorised form uses the identity
which turns the distance table into one matrix product X @ C.T of shape (n,k) plus two norm terms that numpy stretches (broadcasts) from shape (n,1) and (1,k) to (n,k) when added, with no Python loop over points. Rounding can leave tiny negative values, so the code clips at zero.
3. Update step. Each centroid becomes the mean of its members. Averaging is the right move because the mean minimises the sum of squared distances to a set of points.
4. Stop when the largest distance any centroid moved is at most tol, or after max_iter iterations, whichever comes first. The max_iter cap protects against slow final creeping; tol says movement that small does not matter.
5. Empty clusters. If no point is nearest to a centroid, its mean is undefined (this happens, for example, when a centroid starts far from all the data, or when two centroids start at the same place: every distance ties, and argmin returns the first of the tied columns, so the second centroid gets no points). The code re-seeds that centroid at the worst-fit point, meaning the data point with the largest squared distance to its own nearest centroid, one common repair (others re-seed at a random point).
One round by hand on six points
Before the full implementation, here is one assignment step and one update step on six points, using the identity above. Point 0 is (0,0) and centroid 0 is (0,2), so the squared distance is 02+22=4 directly, and by the identity it is 0−0+4=4.
import numpy as np
# six points in 2-D and two starting centroids
X = np.array([[0, 0], [0, 2], [1, 1], [8, 8], [9, 7], [8, 10]], dtype=float)
C = np.array([[0, 2], [9, 7]], dtype=float)
# squared distance of point 0 = (0, 0) to centroid 0 = (0, 2), both ways
direct = ((X[0] - C[0]) ** 2).sum()
expanded = (X[0] ** 2).sum() - 2 * X[0] @ C[0] + (C[0] ** 2).sum()
print("point 0 to centroid 0, direct:", direct, " expanded:", expanded)
d2 = (X ** 2).sum(axis=1, keepdims=True) - 2 * X @ C.T + (C ** 2).sum(axis=1)[None, :]
print("squared distances (rows = points, columns = centroids):")
print(d2)
labels = d2.argmin(axis=1)
print("assignment:", labels.tolist())
print("updated centroids:")
print(np.round(np.array([X[labels == j].mean(axis=0) for j in range(2)]), 3))
# two identical centroids: every distance ties, and argmin takes the first column
C_dup = np.array([[0, 2], [0, 2]], dtype=float)
d2_dup = (X ** 2).sum(axis=1, keepdims=True) - 2 * X @ C_dup.T + (C_dup ** 2).sum(axis=1)[None, :]
print("duplicate centroids, assignment:", d2_dup.argmin(axis=1).tolist())
Output:
point 0 to centroid 0, direct: 4.0 expanded: 4.0
squared distances (rows = points, columns = centroids):
[[ 4. 130.]
[ 0. 106.]
[ 2. 100.]
[100. 2.]
[106. 0.]
[128. 10.]]
assignment: [0, 0, 0, 1, 1, 1]
updated centroids:
[[0.333 1. ]
[8.333 8.333]]
duplicate centroids, assignment: [0, 0, 0, 0, 0, 0]
Read each row of the distance matrix as one point: point 3 is at squared distance 100 from centroid 0 and 2 from centroid 1, so argmin picks column 1. The first three points go to centroid 0 and the last three to centroid 1. The update step then averages each group: the first centroid moves to (1/3,1) and the second to (25/3,25/3), printed as 0.333 and 8.333. The same column gives the k-means++ sampling step: if centroid 0 were the only centroid chosen so far, the next centroid would be drawn with probability proportional to its squared distances 4, 0, 2, 100, 106, 128, which sum to 340, so point 5 is picked with probability 128/340 = 0.376 and the already-chosen point 1 (distance 0) can never be picked again. With the duplicate pair of centroids, every point ties and all six go to the first one, which is the empty-cluster case the code below repairs.
Code, with a check against scikit-learn
The check starts both implementations from the same initial centroids, with n_init=1, algorithm="lloyd" and tol=0.0, so there is no randomness left to differ. It then prints inertia over 30 seeds for each init method with single starts, exercises the input guard, and exercises the empty-cluster repair with duplicate starting centroids and with a starting centroid placed far from all the data.
import numpy as np
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs
def sq_dists(X, C):
"""Squared distance from every row of X (n, d) to every row of C (k, d) -> (n, k)."""
x2 = (X ** 2).sum(axis=1, keepdims=True) # (n, 1)
c2 = (C ** 2).sum(axis=1)[None, :] # (1, k)
d2 = x2 - 2.0 * X @ C.T + c2 # (n, k)
return np.maximum(d2, 0.0) # clip tiny negatives from rounding
def init_centroids(X, k, method, rng):
n = len(X)
if method == "random":
return X[rng.choice(n, size=k, replace=False)].copy()
if method == "kmeans++":
C = [X[rng.integers(n)]]
for _ in range(k - 1):
d2 = sq_dists(X, np.array(C)).min(axis=1) # distance to nearest chosen centre
C.append(X[rng.choice(n, p=d2 / d2.sum())]) # sample proportional to d2
return np.array(C)
raise ValueError(f"unknown init: {method}")
def kmeans(X, k, init="kmeans++", max_iter=100, tol=1e-6, seed=0, C0=None):
X = np.asarray(X, dtype=float)
if X.ndim != 2 or not (1 <= k <= len(X)):
raise ValueError("X must be 2-D and 1 <= k <= n_samples")
rng = np.random.default_rng(seed)
C = init_centroids(X, k, init, rng) if C0 is None else np.array(C0, dtype=float)
for it in range(1, max_iter + 1):
d2 = sq_dists(X, C)
labels = d2.argmin(axis=1) # assignment step
newC = C.copy()
for j in range(k): # update step
members = X[labels == j]
if len(members):
newC[j] = members.mean(axis=0)
else: # empty cluster: re-seed at the worst-fit point
newC[j] = X[d2.min(axis=1).argmax()]
shift = np.sqrt(((newC - C) ** 2).sum(axis=1)).max()
C = newC
if shift <= tol:
break
d2 = sq_dists(X, C)
labels = d2.argmin(axis=1)
inertia = d2[np.arange(len(X)), labels].sum()
return C, labels, inertia, it
X, blob = make_blobs(n_samples=300, centers=4, cluster_std=1.0, random_state=1)
# Same starting centroids for both implementations, so the comparison is exact
C0 = init_centroids(X, 4, "kmeans++", np.random.default_rng(0))
C, labels, inertia, iters = kmeans(X, 4, C0=C0, tol=0.0, max_iter=300)
ref = KMeans(n_clusters=4, init=C0, n_init=1, algorithm="lloyd", tol=0.0, max_iter=300).fit(X)
print("iterations used:", iters)
print("centroids match sklearn:", np.allclose(C, ref.cluster_centers_))
print("labels match sklearn:", bool((labels == ref.labels_).all()))
print("inertia mine vs sklearn:", round(inertia, 4), round(ref.inertia_, 4))
# Effect of the init choice over 30 seeds, single start each
for method in ["random", "kmeans++"]:
runs = [kmeans(X, 4, init=method, seed=s) for s in range(30)]
inertias = [r[2] for r in runs]
print(f"{method:9s} single-start inertia: best {min(inertias):.1f}, mean {np.mean(inertias):.1f}, "
f"worst {max(inertias):.1f}, runs above best+1: {sum(i > min(inertias) + 1 for i in inertias)} of 30")
worst = runs[int(np.argmax(inertias))][1]
print(" worst run, true blob of each member per cluster:", [np.bincount(blob[worst == j], minlength=4).tolist() for j in range(4)])
# Guard checks
for bad_k in (0, 301):
try:
kmeans(X, bad_k)
except ValueError as e:
print("k =", bad_k, "->", e)
# Empty-cluster handling: two identical starting centroids leave one cluster empty on iteration 1
C0_dup = np.vstack([C0[0], C0[0], C0[1], C0[2]])
_, lab_dup, _, _ = kmeans(X, 4, C0=C0_dup)
print("cluster sizes after duplicate start:", np.bincount(lab_dup, minlength=4).tolist())
# A starting centroid far from all the data would stay empty forever without the re-seed
C0_far = C0.copy()
C0_far[3] = [1000.0, 1000.0]
_, lab_far, _, _ = kmeans(X, 4, C0=C0_far)
sizes_far = np.bincount(lab_far, minlength=4)
print("cluster sizes after a far-away start:", sizes_far.tolist())
assert sizes_far.min() > 0
Output:
iterations used: 5
centroids match sklearn: True
labels match sklearn: True
inertia mine vs sklearn: 575.4941 575.4941
random single-start inertia: best 575.5, mean 689.4, worst 1068.2, runs above best+1: 7 of 30
worst run, true blob of each member per cluster: [[39, 0, 0, 0], [0, 1, 75, 2], [0, 74, 0, 73], [36, 0, 0, 0]]
kmeans++ single-start inertia: best 575.5, mean 625.2, worst 1083.9, runs above best+1: 3 of 30
worst run, true blob of each member per cluster: [[35, 0, 0, 0], [0, 1, 75, 2], [0, 74, 0, 73], [40, 0, 0, 0]]
k = 0 -> X must be 2-D and 1 <= k <= n_samples
k = 301 -> X must be 2-D and 1 <= k <= n_samples
cluster sizes after duplicate start: [75, 77, 75, 73]
cluster sizes after a far-away start: [75, 75, 73, 77]
Reading the output: centroids, labels and inertia match scikit-learn exactly from the same start (inertia 575.4941 both), converging in 5 iterations. Over 30 single starts on this one dataset, the best inertia found by either init was 575.5. 7 random starts and 3 k-means++ starts finished more than 1 above that best, and the mean inertia is lower for k-means++ (see the output). That matches the usual reason k-means++ is the default: it makes unlucky starts less frequent. It does not remove them, and both methods still have a worst run above 1000, so you still keep the best of several starts. The worst runs show what an unlucky start looks like: the last line of each init prints, for each cluster, how many members come from each of the four real blobs, and in both worst runs one real blob has been cut into two clusters while two other blobs share a single cluster. This is one dataset and 30 seeds, which shows the pattern but is not a ranking of the two methods in general. With the duplicate start the repair fires on the first iteration and all four clusters end non-empty (sizes 75, 77, 75, 73), but that case alone does not prove the repair is needed: once the first update has moved the other centroid, the stale copy is no longer a duplicate and could pick up points on its own. The far-away start is the case that needs it. A centroid placed at (1000, 1000) is never the nearest centroid to any point, so without the re-seed it would stay empty for every iteration; with it, all four clusters end non-empty (sizes 75, 75, 73, 77), and the assert in the code fails if the re-seed line is removed.
Complexity and memory
Each iteration costs O(nkd) time (the XC⊤ product dominates) and O(nk) memory for the distance matrix, plus O(nd) for the data. For very large n with large k, process X in row blocks so the (n,k) matrix never exists at once. The number of iterations is usually small but has no tight general bound.
Edge cases and pitfalls
- Local optimum: run several starts (different seeds) and keep the lowest inertia; this is what
n_initdoes in libraries. - Empty clusters: handled above; otherwise you divide by zero or produce NaN centroids.
- k outside 1 to n: the guard rejects it. This version assumes the data holds at least k distinct points.
- Feature scale: distances are Euclidean, so a column in dollars swamps a column in counts. Standardise first.
- Non-numeric or non-finite input: this version does not check for NaN or infinity; a production version must reject them before the loop, because a single NaN poisons every distance.
- Convergence test: testing that the assignments stopped changing is an alternative to the centroid-shift test used here.
Unlock Full Question Bank
Get access to all 14 Predictive Modeling and Machine Learning Fundamentals interview questions and detailed answers.
Sign in to ContinueJoin thousands of developers preparing for their dream job.