numpy 的一些小技巧

在此记录一些关于 Numpy 使用的技巧

1. 快速筛选出数组中满足要求的数

x = np.arange(1, 10)  # array([1, 2, 3, 4, 5, 6, 7, 8, 9])
x[(x>1)&(x<6)&(x%2==1)]   # 输出array([3, 5])

numpy贴心地为我们重载好了布尔数组的&运算符，多个条件也不怕。以前我都是用x[np.where(x>1)]，但是这样一次只能筛选一个条件。

2. 找到数组中某个特定数的索引

a = np.array([1, 2, 3])
np.argwhere(a==1)   # 输出 array([[0]])

np.argwhere这个函数的本意是输出数组中非零值的索引，用在这里刚刚好能解决问题。np.argmax只返回一个最大数的索引，因此可以用np.argwhere(a==a.max()) 来找到数组最大数的全部索引。

3. 连接数组

a = np.array([1, 2, 3])
b = np.array([4, 5, 6])

np.r_[a, b, a]
# array([1, 2, 3, 4, 5, 6, 1, 2, 3])
np.concatenate([a, b, a])
# array([1, 2, 3, 4, 5, 6, 1, 2, 3])

np.r_提供了非常简洁的拼接行向量的方法！与之类似的还有np.c_。

4. 压缩/添加维度

x = np.array([[[0], [1], [2]]])
x.shape
# (1, 3, 1)
np.squeeze(x).shape
# (3,)

np.squeeze把数组多余的维度给“挤压”掉，而np.unsqueeze用于添加一个维度。