11. ML Coding 基础

def euc_dist(value1, value2):
    # implement this
    pass

def solution(train_data, test_data, k):
    final_labels = list()
    ...                          # 题目给好的胶水代码：调用辅助函数、多数投票，不要改
    return final_labels

import numpy as np

a, b = np.array([1, 2, 3]), np.array([1, 2.7, 3])   # (3,)、(3,)
print(a.dtype, b.dtype, b.astype(int))       # int64 float64 [1 2 3]，astype 向 0 截断，不四舍五入
print(np.zeros((2, 3)).shape, np.eye(2).tolist())   # (2, 3) [[1.0, 0.0], [0.0, 1.0]]，同类有 np.ones、np.full
print(np.arange(0, 1, 0.25), np.linspace(0, 1, 3))  # [0.   0.25 0.5  0.75] [0.  0.5 1. ]，arange 不含终点
probs, y = np.array([[0.1, 0.9], [0.8, 0.2]]), np.array([1.0, 0.0])   # (N, C)、(N,) 的 float 标签
print(probs[np.arange(2), y.astype(int)])    # [0.9 0.8]，不转 int 直接当下标会报 IndexError
rng = np.random.default_rng(0)
print(rng.integers(0, 10, size=5), rng.normal(0.0, 1.0, size=(2, 3)).shape)   # [8 6 5 2 3] (2, 3)
print(rng.choice(5, size=3, replace=False), rng.permutation(5))   # [4 2 0] [2 1 3 4 0]，同一个 perm 打乱 X 和 y
print(rng.choice(3, size=5, p=[0.1, 0.2, 0.7]))   # [0 2 1 2 2]，按概率抽，K-Means++ 和采样解码都用它

X = np.arange(6).reshape(2, 3)               # (2, 3)，按行填入：[[0 1 2] [3 4 5]]
print(X.reshape(3, -1).tolist(), X.T.tolist())   # [[0, 1], [2, 3], [4, 5]] [[0, 3], [1, 4], [2, 5]]，reshape 不是转置
x = np.arange(6)                             # (6,)
print(x[:, None].shape, x[None, :].shape, x[:, None].squeeze().shape)   # (6, 1) (1, 6) (6,)，None 插入长度 1 的轴
print(np.zeros((4, 2, 3)).transpose(0, 2, 1).shape)   # (4, 3, 2)，相当于 torch 的 transpose(-2, -1)
print(np.concatenate([X, X], axis=0).shape, np.concatenate([X, X], axis=1).shape)   # (4, 3) (2, 6)，沿已有的轴拼
print(np.stack([X, X]).shape, np.hstack([np.ones((2, 1)), X]).shape)   # (2, 2, 3) (2, 4)，后者是加一列 1 当截距

X = np.arange(12).reshape(3, 4)              # (3, 4)：[[0 1 2 3] [4 5 6 7] [8 9 10 11]]
print(X[1:, ::2].tolist())                   # [[4, 6], [8, 10]]
print(X[:, 0].shape, X[:, [0]].shape)        # (3,) (3, 1)，整数下标去掉这一维，列表保留
mask = (X[:, 0] > 3) & ~(X[:, 1] == 9)       # (3,)
print(mask, X[mask].shape, X[X % 5 == 0])    # [False  True False] (1, 4) [ 0  5 10]，二维掩码取出来是一维
y = np.array([2, 0, 3])                      # (3,)，每行要取的列号
print(X[np.arange(3), y])                    # [ 2  4 11]，第 i 行取第 y[i] 列，交叉熵取真实类别就这么写
print(X[:, y].shape)                         # (3, 3)，常见错写法：每行都取这 3 列，不报错
row = X[0, :2]                               # 切片是 view
row[0] = -1
print(X[0])                                  # [-1  1  2  3]，原数组也被改了

X = np.array([[1, 2, 3], [4, 5, 6]])         # (2, 3)
print(X.sum(axis=0), X.sum(axis=1), X.sum(axis=1, keepdims=True).shape)   # [5 7 9] [ 6 15] (2, 1)
print((X + np.array([10, 20, 30])).shape)    # (2, 3)，(2, 3) + (3,)：每行加同一个向量
y = np.array([1.0, 2.0, 3.0])                # (3,)
y_pred = np.array([[1.1], [1.9], [3.2]])     # (3, 1)，模型输出常见的形状
print((y_pred - y).shape)                    # (3, 3)，静默变成两两相减
print(np.mean((y_pred - y) ** 2).round(4), np.mean((y_pred.reshape(-1) - y) ** 2).round(4))   # 1.42 0.02
S = np.arange(1.0, 10.0).reshape(3, 3)       # (3, 3) 方阵，比如 (T, T) 的注意力分数
print((S / S.sum(axis=1)).sum(axis=1))       # [0.425 1.25  2.075]，漏了 keepdims：不报错，行和却不是 1
print((S / S.sum(axis=1, keepdims=True)).sum(axis=1))   # [1. 1. 1.]

x = np.array([3, 1, 4, 1, 5])                # (5,)
print(x.sum(), x.mean(), x.max(), x.argmax())   # 14 2.8 5 4，argmax 并列时取第一个
print(x.std(), x.std(ddof=1), np.cumsum(x))  # 1.6 1.7888543819998317 [ 3  4  8  9 14]
logits = np.array([[0.1, 2.0, -1.0], [1.5, 0.3, 0.2]])   # (N, C)
print((logits.argmax(axis=1) == np.array([1, 2])).mean())   # 0.5，准确率
scores = np.array([0.2, 0.9, 0.1, 0.7, 0.5])   # (5,)
top = np.argpartition(scores, -2)[-2:]       # (2,)，最大的 2 个，顺序不保证
print(np.argsort(-scores), top[np.argsort(-scores[top])])   # [1 3 4 0 2] [1 3]
D = np.array([[3, 1, 2], [0, 5, 4]])         # (2, 3)，2 个查询点到 3 个样本的距离
idx = np.argsort(D, axis=1)[:, :2]           # (2, 2)，每行最近的 2 个
print(np.take_along_axis(D, idx, axis=1).tolist())   # [[1, 2], [0, 4]]，逐行按下标取值
values, counts = np.unique(np.array([2, 0, 2, 1, 2]), return_counts=True)
print(values, counts, values[counts.argmax()])   # [0 1 2] [1 1 3] 2，KNN 多数投票

x = np.array([-2.0, -0.5, 0.0, 1.5, 3.0])   # (5,)
print(np.where(x > 0, x, 0.0), np.where(x > 0)[0])   # [0.  0.  0.  1.5 3. ] [3 4]，只传条件时返回下标
print(np.clip(x, -1.0, 2.0))                 # [-1.  -0.5  0.   1.5  2. ]，取 log 前写 np.clip(p, 1e-15, 1 - 1e-15)
print(np.maximum(x, 0), np.max(x))           # [0.  0.  0.  1.5 3. ] 3.0
A, B = np.array([[1, 2], [3, 4]]), np.array([[5, 6], [7, 8]])   # (2, 2)、(2, 2)
print((A * B).tolist(), (A @ B).tolist())    # [[5, 12], [21, 32]] [[19, 22], [43, 50]]
print(np.linalg.norm(np.array([[3.0, 4.0], [6.0, 8.0]]), axis=1))   # [ 5. 10.]，每行的 L2 范数
print(np.einsum("bqd,bkd->bqk", np.ones((2, 4, 8)), np.ones((2, 5, 8))).shape)   # (2, 4, 5)，即 Q @ K.transpose(0, 2, 1)

from typing import List


def solution(matrix: List[List[float]]) -> List[List[float]]:
    """按列 z-score，常数列输出 0。matrix: (n, d) -> (n, d)"""
    if not matrix:
        return []
    X = np.asarray(matrix, dtype=float)                  # 1. (n, d)，转 float 防整型截断
    const = X.max(axis=0) == X.min(axis=0)               # 2. (d,)，常数列（不用 std == 0，见 0.4 节）
    std = np.where(const, 1.0, X.std(axis=0))            # 3. (d,)，常数列分母换成 1
    return np.where(const, 0.0, (X - X.mean(axis=0)) / std).tolist()   # 4. (n, d)，转回 Python 类型


if __name__ == "__main__":
    out = solution([[1, 10, 5], [2, 20, 5], [3, 30, 5]])
    assert np.allclose(out, [[-1.2247, -1.2247, 0], [0, 0, 0], [1.2247, 1.2247, 0]], atol=1e-4)
    assert solution([[0.1], [0.1], [0.1]]) == [[0.0], [0.0], [0.0]] and solution([]) == []
    print("all tests passed")

import io

import numpy as np
import pandas as pd

df = pd.DataFrame({
    "user_id": [1, 2, 3, 4, 5, 6, 7, 8],
    "city": ["NY", "SF", "NY", "LA", "SF", "NY", "NY", "SF"],
    "age": [25, 32, np.nan, np.nan, 28, 35, 45, np.nan],
    "income": [50.0, 120.0, 65.0, np.nan, 95.0, 70.0, 80.0, 110.0],
    "bought": [0, 1, 0, 1, 1, 0, 1, 0],
})
print(df.shape, df.dtypes["age"], df.dtypes["city"])   # (8, 5) float64 object，有 NaN 的整数列变成 float
print(df.describe().loc["mean"].round(1).to_dict())    # {'user_id': 4.5, 'age': 33.0, 'income': 84.3, 'bought': 0.5}
print(df["city"].value_counts().to_dict())             # {'NY': 4, 'SF': 3, 'LA': 1}，再看 df.head()、df.info()
print(pd.read_csv(io.StringIO("a,b\n1,\n2,x\n")).isna().sum().to_dict())   # {'a': 0, 'b': 1}，空字段读成 NaN

age, sub = df["age"], df[["city", "age"]]       # 单括号是 Series (8,)，双括号是 DataFrame (8, 2)
print(df.loc[2, "city"], df.iloc[2, 1], len(df.loc[1:3]), len(df.iloc[1:3]))   # NY NY 3 2，loc 按标签且含右端
mask = (df["city"] == "NY") & (df["age"] > 30)  # 每个条件加括号，用 & | ~；和 NaN 比较一律 False
print(df.loc[mask, "user_id"].tolist())         # [6, 7]，一次 .loc 同时过滤行、选列
print(df.loc[~df["city"].isin(["SF", "LA"]), "user_id"].tolist())   # [1, 3, 6, 7]
print(df.loc[df["income"].between(60, 95), "user_id"].tolist())     # [3, 5, 6, 7]，两端都含

d = df.copy()
d["income_k"] = d["income"] / 1000                          # 整列运算，NaN 自动传播
d["is_senior"] = np.where(d["age"] >= 30, 1, 0)             # 向量化 if-else；NaN >= 30 是 False
d["region"] = d["city"].map({"NY": "east", "SF": "west"})   # 字典里没有的键变成 NaN
print(d["is_senior"].tolist(), d["region"].tolist()[:4])    # [0, 1, 0, 0, 0, 1, 1, 0] ['east', 'west', 'east', nan]
print(df.isna().sum().to_dict())   # {'user_id': 0, 'city': 0, 'age': 3, 'income': 1, 'bought': 0}
d["income"] = d["income"].fillna(d["income"].median())      # 整列填中位数 80.0
d["age"] = d["age"].fillna(d.groupby("city")["age"].transform("mean"))   # 组内填充：所在城市的平均年龄
d["age"] = d["age"].fillna(df["age"].median())              # LA 整组缺失，再用全局中位数兜底
print(d["age"].tolist())           # [25.0, 32.0, 35.0, 32.0, 28.0, 35.0, 45.0, 30.0]
print(df.dropna().shape, df.dropna(subset=["income"]).shape)   # (5, 5) (7, 5)

stats = df.groupby("city").agg(n_users=("user_id", "count"), avg_income=("income", "mean"),
                               buy_rate=("bought", "mean"))      # 新列名=(原列名, 聚合函数字符串)
print(stats.round(2))
#       n_users  avg_income  buy_rate
# city
# LA          1         NaN      1.00
# NY          4       66.25      0.25
# SF          3      108.33      0.67
diff = df["income"] - df.groupby("city")["income"].transform("mean")   # (8,)，减所在城市的均值
print(diff.round(2).tolist())      # [-16.25, 11.67, -1.25, nan, -13.33, 3.75, 13.75, 1.67]
by_income = df.sort_values("income", ascending=False)                  # NaN 默认排最后
print(by_income["user_id"].tolist(), df.nlargest(3, "income")["user_id"].tolist())   # [2, 8, 5, 7, 6, 3, 1, 4] [2, 8, 5]
print(by_income.groupby("city").head(2)["user_id"].tolist())           # [2, 8, 7, 6, 4]，每个城市收入前 2

users = df[["user_id", "city"]]
orders = pd.DataFrame({"user_id": [1, 1, 2, 5, 9], "amount": [30.0, 20.0, 50.0, 15.0, 99.0]})
print([users.merge(orders, on="user_id", how=h).shape for h in ["inner", "left", "outer"]])   # [(4, 3), (9, 3), (10, 3)]
spend = orders.groupby("user_id", as_index=False).agg(total=("amount", "sum"))
feat = users.merge(spend, on="user_id", how="left", validate="one_to_one")
print(feat["total"].fillna(0.0).tolist())      # [50.0, 50.0, 0.0, 0.0, 15.0, 0.0, 0.0, 0.0]
print(pd.concat([users, users], ignore_index=True).shape, pd.concat([users, df[["age"]]], axis=1).shape)   # (16, 2) (8, 3)
table = df.pivot_table(index="city", columns="bought", values="user_id", aggfunc="count", fill_value=0)
print(table.values.tolist())                   # [[0, 1], [3, 1], [1, 2]]，行是 LA/NY/SF，列是 bought=0/1
X = df[["age", "income"]].fillna(0.0).to_numpy(dtype=np.float32)   # (8, 2)，接着 torch.from_numpy(X)

from typing import Tuple

NUM, CAT = ["age", "income"], ["city"]
TRAIN_CSV = "age,income,city,bought\n25,50,NY,0\n32,120,SF,1\n,65,NY,0\n41,,LA,1\n28,95,SF,1\n35,70,NY,0\n"
TEST_CSV = "age,income,city,bought\n30,,NY,0\n,100,SF,1\n52,85,Austin,1\n"


def preprocess(train: pd.DataFrame, test: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """所有统计量只用 train。(n_train, d)、(n_test, d) -> 列名和顺序相同的两张表"""
    fill = train[NUM].median().to_dict()                     # 1. 数值列填 train 的中位数
    fill.update({c: "unknown" for c in CAT})                 #    类别列填 "unknown"
    train, test = train.fillna(fill), test.fillna(fill)
    train = pd.get_dummies(train, columns=CAT, dtype=int)    # 2. one-hot，test 缺的列补 0、多的列丢掉
    test = pd.get_dummies(test, columns=CAT, dtype=int).reindex(columns=train.columns, fill_value=0)
    mean, std = train[NUM].mean(), train[NUM].std(ddof=0)    # 3. z-score，ddof=0 和 StandardScaler 一致
    std = std.where(std > 1e-8, 1.0)                         #    常数列除以 1
    train[NUM], test[NUM] = (train[NUM] - mean) / std, (test[NUM] - mean) / std
    return train, test


if __name__ == "__main__":
    train, test = pd.read_csv(io.StringIO(TRAIN_CSV)), pd.read_csv(io.StringIO(TEST_CSV))   # 考试里读 data/ 下的文件
    X_tr, X_te = preprocess(train.drop(columns=["bought"]), test.drop(columns=["bought"]))
    assert list(X_tr.columns) == list(X_te.columns) and X_te.notna().all().all()
    assert np.allclose(X_tr[NUM].mean(), 0.0) and np.allclose(X_tr[NUM].std(ddof=0), 1.0)
    print(X_te.round(2))
    print("all tests passed")
#     age  income  city_LA  city_NY  city_SF
# 0 -0.43   -0.36        0        1        0
# 1 -0.03    0.95        0        0        1
# 2  3.90    0.29        0        0        0      <- Austin 没见过，三个 city 列都是 0
# all tests passed

import numpy as np
import torch

a, b = torch.tensor([1, 2, 3]), torch.tensor([1.0, 2.0])    # 小写 tensor：从数据建，自动推断 dtype
c = torch.from_numpy(np.array([1.0, 2.0]))                  # 和 NumPy 共享内存，还是 float64
print(a.dtype, b.dtype, c.dtype, torch.zeros(2, 3).shape)   # torch.int64 torch.float32 torch.float64 torch.Size([2, 3])
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x32 = c.float().to(device)                    # x.to() 返回新张量要接住；model.to(device) 原地移动
print(x32.dtype, x32.device, b.long().dtype)  # torch.float32 cpu torch.int64，标签和 token id 一律 long

x = torch.arange(6).reshape(2, 3)             # (2, 3)，一律用 reshape：内存不连续时它自动复制
print(x.t().is_contiguous(), x.t().reshape(6))    # False tensor([0, 3, 1, 4, 2, 5])，这里写 view(6) 会报错：view 只认连续内存
B, T, H, D_h = 2, 5, 4, 8
q = torch.randn(B, T, H * D_h).reshape(B, T, H, D_h).transpose(1, 2)   # (B, H, T, D_h)：拆多头
merged = q.transpose(1, 2).reshape(B, T, H * D_h)                      # (B, T, C)：拼回多头；写 view 的话要先 .contiguous()
print(q.permute(0, 2, 1, 3).shape, merged.shape)   # torch.Size([2, 5, 4, 8]) torch.Size([2, 5, 32])
print(x[0].unsqueeze(0).shape, x[0].unsqueeze(1).shape)   # torch.Size([1, 3]) torch.Size([3, 1])
print(torch.zeros(1, 3, 1).squeeze(-1).shape) # torch.Size([1, 3])，squeeze 总带 dim，免得 batch=1 也被挤掉
print(torch.cat([x, x], dim=0).shape, torch.stack([x, x], dim=0).shape)   # torch.Size([4, 3]) torch.Size([2, 2, 3])

logits = torch.tensor([[2.0, 1.0, 0.1], [0.5, 2.5, 0.3]])   # (N, C)
y = torch.tensor([0, 1])                      # (N,)，long；gather 的 index 和输出同形状，先 unsqueeze 成 (N, 1)
print(logits[torch.arange(2), y], logits.gather(1, y.unsqueeze(1)).squeeze(1))   # 两个都是 tensor([2.0000, 2.5000])
print(logits.topk(2, dim=-1).indices.tolist(), logits.argmax(dim=-1))   # [[0, 1], [1, 0]] tensor([0, 1])
print(logits.sum(dim=1, keepdim=True).shape, logits.max(dim=1).values)  # torch.Size([2, 1]) tensor([2.0000, 2.5000])
causal = torch.tril(torch.ones(3, 3))         # (T, T) 下三角：1 = 可见，0 = 屏蔽；masked_fill 返回新张量
print(torch.zeros(3, 3).masked_fill(causal == 0, -1e9).softmax(dim=-1)[1])   # tensor([0.5000, 0.5000, 0.0000])
A = torch.randn(2, 5, 8)                      # (B, T, C)
print((A @ torch.randn(8, 16)).shape, torch.bmm(A, torch.randn(2, 8, 7)).shape)   # torch.Size([2, 5, 16]) torch.Size([2, 5, 7])，bmm 只收 3 维

w = torch.tensor([1.0, 2.0], requires_grad=True)   # 叶子张量
(w ** 2).sum().backward()                     # backward 只能对标量调用
(w ** 2).sum().backward()
print(w.grad)                                 # tensor([4., 8.])，两次累加：2w + 2w
with torch.no_grad():
    w -= 0.1 * w.grad                         # 手动 SGD 一步；不包 no_grad 会报错
w.grad.zero_()                                # optimizer.zero_grad() 做的就是这个
print(w, (w * 3).detach().requires_grad, round((w ** 2).sum().item(), 4))   # tensor([0.6000, 1.2000], requires_grad=True) False 1.8

import torch.nn as nn
import torch.nn.functional as F


class MLP(nn.Module):
    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int = 2):
        super().__init__()                             # 1. 先调父类构造：漏写的话，下一行赋值子模块就报 AttributeError
        self.inp = nn.Linear(in_dim, hidden_dim)       # 2. 有参数的层写成属性才会注册
        layers = [nn.Linear(hidden_dim, hidden_dim) for _ in range(num_layers)]
        self.hidden = nn.ModuleList(layers)            # 3. 必须包 ModuleList：普通 list 不注册，优化器看不到，也不报错
        self.dropout = nn.Dropout(0.1)                 # 自己会看 train() / eval()
        self.out = nn.Linear(hidden_dim, out_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """x: (B, in_dim) -> logits: (B, out_dim)，不过 softmax"""
        h = F.relu(self.inp(x))                        # (B, hidden_dim)
        for layer in self.hidden:
            h = self.dropout(F.relu(layer(h)))         # (B, hidden_dim)
        return self.out(h)                             # (B, out_dim)


model = MLP(4, 8, 3)
print(model(torch.randn(5, 4)).shape, sum(p.numel() for p in model.parameters()))   # torch.Size([5, 3]) 211

import io
from typing import Tuple

from torch.utils.data import DataLoader, Dataset


class ToyDataset(Dataset):
    """dataset[i] = (X[i], y[i])，和现成的 TensorDataset(X, y) 等价"""
    def __init__(self, X: torch.Tensor, y: torch.Tensor):
        self.X, self.y = X, y                         # (N, D)、(N,)

    def __len__(self) -> int:
        return len(self.X)

    def __getitem__(self, i: int) -> Tuple[torch.Tensor, torch.Tensor]:
        return self.X[i], self.y[i]


torch.manual_seed(42)
X = torch.randn(64, 4)                                # (N, D)
y = (X[:, 0] + X[:, 1] > 0).long()                    # (N,)，long 标签
loader = DataLoader(ToyDataset(X, y), batch_size=16, shuffle=True)
model, criterion = MLP(4, 16, 2), nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(20):                               # 完整的训练、验证循环和早停见专题 10.1 节
    model.train()                                     # 1. Dropout 生效
    for xb, yb in loader:                             # xb: (16, 4)，yb: (16,)
        optimizer.zero_grad()                         # 2. 清梯度（2.x 默认把 grad 设成 None）
        loss = criterion(model(xb), yb)               # 3. 前向，logits (16, 2)
        loss.backward()                               # 4. 反向
        optimizer.step()                              # 5. 更新
model.eval()                                          # 6. 关掉 Dropout
with torch.no_grad():                                 # 7. 不建计算图
    print((model(X).argmax(dim=1) == y).float().mean().item())   # 1.0，train acc
buffer = io.BytesIO()                                 # 真实场景写 "model.pt"
torch.save(model.state_dict(), buffer)                # 只存「名字 -> 张量」字典
buffer.seek(0)
print(MLP(4, 16, 2).load_state_dict(torch.load(buffer, weights_only=True)))   # <All keys matched successfully>

def get_shape(A: list) -> tuple:
    """返回 (行数, 列数)；空矩阵记为 (0, 0)；每一行必须一样长"""
    n_cols = len(A[0]) if A else 0
    if any(len(row) != n_cols for row in A):
        raise ValueError(f"ragged matrix: every row must have {n_cols} columns, like row 0")
    return len(A), n_cols

def transpose(A: list) -> list:
    """A: (n, m) -> (m, n)。一行写法 [list(col) for col in zip(*A)]：zip 产出 tuple，要转回 list"""
    n, m = get_shape(A)
    return [[A[i][j] for i in range(n)] for j in range(m)]   # 外层走列 j，内层走行 i；非方阵别都写成 range(n)

def matmul(A: list, B: list) -> list:
    """A: (n, m), B: (m, p) -> A @ B: (n, p)"""
    (n, m), (m_b, p) = get_shape(A), get_shape(B)
    if n > 0 and m != m_b:
        raise ValueError(f"shape mismatch: A is ({n}, {m}), B is ({m_b}, {p}); need A cols == B rows")
    C = [[0.0] * p for _ in range(n)]       # (n, p)；不能写 [[0.0] * p] * n，那样 n 行是同一个 list
    for i in range(n):
        for k in range(m):
            a_ik, row_b = A[i][k], B[k]     # 整个内层循环都不变
            for j in range(p):
                C[i][j] += a_ik * row_b[j]
    return C

def solution(A: list, B: list) -> list:
    return matmul(A, B)

def normalize(X: list, method: str = "z-score") -> list:
    """X: (N, D) -> (N, D)，每列做 min-max 或 z-score（总体标准差）；常数列输出 0"""
    if method not in ("z-score", "min-max"):
        raise ValueError(f"unknown method: {method!r}")
    N, cols = len(X), transpose(X)                                  # cols: (D, N)；空矩阵时为 []
    const = [max(c) == min(c) for c in cols]                        # (D,)：用原始数据判断常数列
    if method == "min-max":
        shift = [min(c) for c in cols]                              # (D,)
        scale = [max(c) - lo for c, lo in zip(cols, shift)]         # (D,)
    else:
        shift = [sum(c) / N for c in cols]                          # (D,)：均值
        scale = [(sum((v - mu) ** 2 for v in c) / N) ** 0.5 for c, mu in zip(cols, shift)]   # 两遍法
    return [[0.0 if is_c else (v - s) / d for v, s, d, is_c in zip(row, shift, scale, const)]
            for row in X]                                           # (N, D)

E = 2.718281828459045                              # 自然常数 e，等于 math.e

def add_bias(Y: list, b: list) -> list:
    """Y: (N, D), b: (D,) -> (N, D)。手写广播：同一个 b 加到每一行"""
    N, D = get_shape(Y)
    if N > 0 and D != len(b):
        raise ValueError(f"shape mismatch: Y is ({N}, {D}), b has length {len(b)}")
    return [[y + bj for y, bj in zip(row, b)] for row in Y]

def sigmoid(z: float) -> float:
    # 只对非正数取指数，E ** 大正数会抛 OverflowError
    return 1.0 / (1.0 + E ** (-z)) if z >= 0 else E ** z / (1.0 + E ** z)

def softmax_row(z: list) -> list:
    """z: (C,) -> (C,)，和为 1；先减最大值，最大的指数项是 E ** 0 = 1"""
    z_max = max(z)
    exps = [E ** (v - z_max) for v in z]
    total = sum(exps)
    return [e / total for e in exps]

ACTIVATIONS = {"relu": lambda z: max(z, 0.0), "sigmoid": sigmoid}

def forward(X: list, W1: list, b1: list, W2: list, b2: list, activation: str = "relu") -> list:
    """X: (N, D_in), W1: (D_in, D_h), b1: (D_h,), W2: (D_h, C), b2: (C,) -> (N, C)，每行和为 1"""
    act = ACTIVATIONS[activation]                                        # 未知的激活函数名会抛 KeyError
    H = [[act(v) for v in row] for row in add_bias(matmul(X, W1), b1)]   # (N, D_h)
    logits = add_bias(matmul(H, W2), b2)                                 # (N, C)
    return [softmax_row(row) for row in logits]                          # (N, C)

def to_sparse(A: list) -> dict:
    """稠密 (n, m) -> {i: {j: A[i][j]}}，只存非零元素，全零的行整行不存"""
    return {i: {j: v for j, v in enumerate(row) if v != 0} for i, row in enumerate(A) if any(row)}

def sparse_matmul(A: list, B: list) -> list:
    """A: (n, m), B: (m, p)，大部分元素是 0 -> A @ B: (n, p)"""
    (n, m), (m_b, p) = get_shape(A), get_shape(B)
    if n > 0 and m != m_b:
        raise ValueError(f"shape mismatch: A is ({n}, {m}), B is ({m_b}, {p}); need A cols == B rows")
    A_sp, B_sp = to_sparse(A), to_sparse(B)
    C = [[0.0] * p for _ in range(n)]                  # (n, p)
    for i, row_a in A_sp.items():
        for k, a in row_a.items():                     # A 第 i 行的非零元素
            for j, b in B_sp.get(k, {}).items():       # B 第 k 行的非零元素
                C[i][j] += a * b
    return C

import numpy as np

def close(mine: list, ref: np.ndarray) -> bool:
    return np.shape(mine) == ref.shape and np.allclose(mine, ref)        # 形状和数值都要对

def test_matrix() -> None:
    rng = np.random.default_rng(0)
    for n, m, p in [(1, 1, 1), (2, 3, 4), (5, 1, 3), (1, 4, 1), (4, 4, 4)]:
        A, B = rng.standard_normal((n, m)), rng.standard_normal((m, p))
        A[rng.random((n, m)) < 0.5] = 0.0                                   # 约一半置 0，给稀疏版用
        assert close(matmul(A.tolist(), B.tolist()), A @ B)
        assert close(sparse_matmul(A.tolist(), B.tolist()), A @ B)
        assert transpose(A.tolist()) == A.T.tolist()                        # 必须是 list，不能是 tuple
    X = rng.standard_normal((6, 4))
    assert close(normalize(X.tolist(), "min-max"), (X - X.min(0)) / (X.max(0) - X.min(0)))
    assert close(normalize(X.tolist()), (X - X.mean(0)) / X.std(0))        # np.std 默认除以 N
    assert normalize([[0.1]] * 3) == normalize([[0.1]] * 3, "min-max") == [[0.0], [0.0], [0.0]]
    X, W1, b1, W2, b2 = args = [rng.standard_normal(s) for s in [(5, 3), (3, 4), (4,), (4, 2), (2,)]]
    for name, act in [("relu", lambda z: np.maximum(z, 0.0)), ("sigmoid", lambda z: 1 / (1 + np.exp(-z)))]:
        logits = act(X @ W1 + b1) @ W2 + b2                                  # (5, 2)
        P = np.exp(logits - logits.max(1, keepdims=True))
        assert close(forward(*[a.tolist() for a in args], activation=name), P / P.sum(1, keepdims=True))
    assert softmax_row([1000.0, 0.0]) == [1.0, 0.0] and sigmoid(-1000.0) == 0.0 and sigmoid(1000.0) == 1.0
    assert matmul([], [[1.0]]) == [] and sparse_matmul([], [[1.0]]) == [] and normalize([]) == []
    try:
        matmul([[1.0, 2.0]], [[1.0, 2.0], [3.0]])                           # B 行长不一致
        raise AssertionError("matmul should raise ValueError")
    except ValueError:
        pass
    print("all tests passed")

if __name__ == "__main__":
    test_matrix()