🐍 Python Stats Cheatsheet

# === K-MEANS CLUSTERING ===
# Find natural customer segments

import numpy as np
import pandas as pd
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt

# Sample data: Customer information
data = pd.DataFrame({
    'customer_id': range(1, 21),
    'age': [25, 45, 32, 52, 28, 48, 35, 55, 22, 50, 30, 46, 27, 53, 33, 49, 24, 51, 29, 47],
    'income': [35, 75, 42, 85, 38, 78, 45, 90, 32, 80, 40, 72, 36, 88, 44, 76, 34, 82, 39, 74],  # $1000s
    'purchase_freq': [2, 12, 3, 15, 2, 13, 4, 18, 1, 14, 3, 11, 2, 16, 4, 12, 2, 15, 3, 13]  # per year
})

# Select features for clustering
features = ['age', 'income', 'purchase_freq']
X = data[features]

# === IMPORTANT: Standardize features (required for K-means!) ===
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# === Find optimal number of clusters (Elbow Method) ===
wcss = []  # Within-Cluster Sum of Squares
K_range = range(1, 11)

for k in K_range:
    kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
    kmeans.fit(X_scaled)
    wcss.append(kmeans.inertia_)

print("=== ELBOW PLOT DATA ===")
for k, score in zip(K_range, wcss):
    print(f"k={k}: WCSS={score:.2f}")

# Plot elbow curve (look for the "elbow")
# plt.plot(K_range, wcss, 'bo-')
# plt.xlabel('Number of Clusters (k)')
# plt.ylabel('WCSS')
# plt.title('Elbow Method')
# plt.show()

# === Run K-means with chosen k (let's say k=3) ===
k = 3
kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
data['cluster'] = kmeans.fit_predict(X_scaled)

print(f"\n=== K-MEANS WITH k={k} ===")

# === Cluster centroids (in original scale) ===
centroids_scaled = kmeans.cluster_centers_
centroids = scaler.inverse_transform(centroids_scaled)

print("\nCluster Centroids (average values):")
centroid_df = pd.DataFrame(centroids, columns=features)
centroid_df.index.name = 'Cluster'
print(centroid_df)

# === Cluster sizes ===
print("\n=== CLUSTER SIZES ===")
cluster_counts = data['cluster'].value_counts().sort_index()
print(cluster_counts)

# === Cluster summaries ===
print("\n=== CLUSTER SUMMARIES ===")
for cluster_id in range(k):
    cluster_data = data[data['cluster'] == cluster_id]
    print(f"\nCluster {cluster_id} ({len(cluster_data)} customers):")
    print(cluster_data[features].describe())

# === Interpret clusters ===
print("\n=== CLUSTER INTERPRETATION ===")
for idx, row in centroid_df.iterrows():
    age = "Young" if row['age'] < 35 else "Middle-aged" if row['age'] < 50 else "Mature"
    income = "low-income" if row['income'] < 45 else "mid-income" if row['income'] < 70 else "high-income"
    freq = "infrequent" if row['purchase_freq'] < 5 else "regular" if row['purchase_freq'] < 12 else "frequent"
    print(f"Cluster {idx}: {age}, {income}, {freq} buyers")

# === Assign new customer to cluster ===
new_customer = pd.DataFrame({
    'age': [40],
    'income': [60],
    'purchase_freq': [8]
})
new_customer_scaled = scaler.transform(new_customer)
cluster_assignment = kmeans.predict(new_customer_scaled)[0]
print(f"\n=== NEW CUSTOMER ASSIGNMENT ===")
print(f"New customer assigned to Cluster {cluster_assignment}")