NumPy 数组基础 —— 创建、索引、运算

目标

• 掌握 ndarray 的创建、属性、索引与切片
• 理解向量化运算与 Python 循环的性能差异
• 掌握广播机制

完整代码

import numpy as np

# ============================================================
# 1. 数组创建
# ============================================================

# 从列表创建
a = np.array([1, 2, 3, 4, 5])
print(f"1D 数组: {a}")

b = np.array([[1, 2, 3], [4, 5, 6]])
print(f"2D 数组:\n{b}")

# 特殊数组
zeros = np.zeros((3, 4))           # 全零
ones = np.ones((2, 3))             # 全一
eye = np.eye(3)                     # 单位矩阵
full = np.full((2, 2), 7)          # 填充指定值
print(f"单位矩阵:\n{eye}")

# 序列数组
arange = np.arange(0, 10, 2)      # 步长为2
linspace = np.linspace(0, 1, 5)   # 等分5份
print(f"arange: {arange}")
print(f"linspace: {linspace}")

# 随机数组
rand = np.random.rand(3, 3)        # [0,1) 均匀分布
randn = np.random.randn(3, 3)      # 标准正态分布
randint = np.random.randint(1, 100, size=(3, 3))  # 随机整数

# 设置随机种子（可复现）
rng = np.random.default_rng(seed=42)
print(f"可复现随机数: {rng.random(5)}")

# ============================================================
# 2. 数组属性
# ============================================================
arr = np.random.randn(4, 3, 28, 28)

print(f"\n形状:   {arr.shape}")      # (4, 3, 28, 28)
print(f"维度数: {arr.ndim}")         # 4
print(f"元素总数: {arr.size}")       # 4*3*28*28 = 9408
print(f"数据类型: {arr.dtype}")     # float64
print(f"元素大小: {arr.itemsize} bytes")
print(f"总内存:  {arr.nbytes} bytes ({arr.nbytes / 1024 / 1024:.2f} MB)")

# 数据类型转换
arr_f32 = arr.astype(np.float32)
print(f"float32 内存减半: {arr_f32.nbytes} bytes")

# ============================================================
# 3. 索引与切片
# ============================================================
arr = np.arange(1, 13).reshape(3, 4)
print(f"\n原始数组:\n{arr}")

# 基本切片
print(f"第2行: {arr[1]}")
print(f"第1列: {arr[:, 0]}")
print(f"子矩阵:\n{arr[0:2, 1:3]}")

# 布尔索引
mask = arr > 6
print(f"大于6的元素: {arr[mask]}")

# 花式索引
print(f"取(0,2)行和(1,3)列:\n{arr[[0, 2], :][:, [1, 3]]}")

# np.ix_ —— 笛卡尔积索引
print(f"np.ix_ 笛卡尔积:\n{arr[np.ix_([0, 2], [1, 3])]}")

# ============================================================
# 4. 形状变换
# ============================================================
a = np.arange(12)

print(f"\nreshape 3×4:\n{a.reshape(3, 4)}")
print(f"reshape 2×2×3:\n{a.reshape(2, 2, 3)}")

# -1 表示自动推断
print(f"自动推断行数:\n{a.reshape(-1, 4)}")

# 转置
m = np.arange(6).reshape(2, 3)
print(f"转置:\n{m.T}")

# 添加/删除维度
x = np.array([1, 2, 3])
print(f"np.newaxis 增加列维度: {x[:, np.newaxis].shape}")  # (3, 1)
print(f"expand_dims: {np.expand_dims(x, axis=0).shape}")     # (1, 3)
print(f"squeeze: {np.ones((1,3,1,5)).squeeze().shape}")     # (3, 5)

# 拼接
a = np.array([[1, 2], [3, 4]])
b = np.array([[5, 6]])
print(f"\n垂直拼接:\n{np.vstack([a, b])}")
print(f"水平拼接:\n{np.hstack([a, b.T])}")

# ============================================================
# 5. 向量化运算（性能对比）
# ============================================================
import time

N = 10_000_000
arr = np.random.randn(N)

# 向量化
start = time.time()
result_np = arr * 2 + 1
print(f"\nNumPy 向量化: {time.time() - start:.4f} 秒")

# Python 循环（不要这样做！）
start = time.time()
result_py = [x * 2 + 1 for x in arr]
print(f"Python 列表推导: {time.time() - start:.4f} 秒")

# 逐元素数学运算
a = np.array([1, 2, 3, 4])
print(f"\n加法: {a + 2}")
print(f"乘法: {a * 3}")
print(f"平方: {a ** 2}")
print(f"指数: {np.exp(a)}")
print(f"对数: {np.log(a)}")

# ============================================================
# 6. 广播机制
# ============================================================

# 标量广播
arr = np.arange(6).reshape(2, 3)
print(f"\n标量广播:\n{arr + 10}")

# 向量广播
row = np.array([10, 20, 30])  # shape (3,)
print(f"向量广播（按行）:\n{arr + row}")

col = np.array([[10], [20]])   # shape (2, 1)
print(f"向量广播（按列）:\n{arr + col}")

# 广播规则：从后往前对齐维度，不匹配的维度必须为 1
# (2, 3) + (1, 3) → (2, 3)  ✅
# (2, 3) + (2, 1) → (2, 3)  ✅
# (2, 3) + (2,  ) → (2, 3)  ✅（自动在前面补 1）

# ============================================================
# 7. 聚合操作
# ============================================================
arr = np.random.randn(4, 5)

print(f"\n全局求和: {arr.sum():.4f}")
print(f"全局均值: {arr.mean():.4f}")
print(f"全局最大值: {arr.max():.4f}")
print(f"全局最小值: {arr.min():.4f}")
print(f"标准差: {arr.std():.4f}")

# 按轴聚合
print(f"\n按列求和 (axis=0): {arr.sum(axis=0)}")  # 对每列求和
print(f"按行求和 (axis=1): {arr.sum(axis=1)}")     # 对每行求和
print(f"每列均值: {arr.mean(axis=0)}")
print(f"每行最大值: {arr.max(axis=1)}")

# argmax / argmin
print(f"每列最大值索引: {arr.argmax(axis=0)}")

运行输出示例

NumPy 向量化: 0.0123 秒
Python 列表推导: 3.4567 秒
# 向量化比列表推导快 ~280 倍！

关键要点

概念	说明
ndarray	同构多维数组，所有元素类型相同
shape	元组，描述各维度大小
dtype	数据类型，如 float64、int32
向量化	一次操作作用于整个数组，底层 C 实现
广播	不同形状数组的自动对齐运算
axis=0	沿行方向（对列操作），结果减少第0维
np.newaxis	增加一个大小为1的维度