np.newaxis等方法的使用

1. np.newaxis

官方解释(np.newaxis)：A convenient alias for None, useful for indexing arrays.

import numpy as np

np.newaxis is None
x = np.array([0,1,2])
print(x[np.newaxis,:],x[None,:], sep='\n')
x[np.newaxis,:] == x[None,:]

输出：

[[0 1 2]]
[[0 1 2]]
array([[ True,  True,  True]])

根据官方文档的解释，x[np.newaxis,:]、x[np.newaxis]和x[None]三者是等价的。

x[newaxis, :].shape
(1, 3)
x[newaxis].shape
(1, 3)
x[None].shape
(1, 3)
x[:, newaxis].shape
(3, 1)

所以说np.newaxis就是给None对象起了一个便于理解的别名。至于为什么None对象能用来给np.array创建一个新的维度，还要参看关于numpy文档关于索引的部分（Structural indexing tools）：
To facilitate easy matching of array shapes with expressions and in assignments, the np.newaxis object can be used within array indices to add new dimensions with a size of 1.

2. np.argsort，np.argpartition

np.argsort、np.argpartition对应于np.sort、np.partition，分别返回排序和划分后相应元素的索引。
Usage:
numpy.argsort(a[, axis=-1, kind=None, order=None])

numpy.argpartition(a, kth[, axis=-1, kind=‘introselect’, order=None])
kind是可选的排序/划分的算法参数。

3. 实例: 利用np.newaxis、np.argsort、np.argpartition等函数实现k邻近

Python数据科学手册[1]实例：求随机点k邻近并可视化。

# np.newaxis 的使用

import numpy as np

x = np.random.rand(10, 2)
# print(x, end='\n\n')
# print(x[:,np.newaxis,:], x[np.newaxis,:,:], sep='\n'*2)
print(x[:,np.newaxis,:].shape, x[np.newaxis,:,:].shape, sep='\n')

# 利用numpy的广播特性 shape(10,1(np.newaxis),2) - shape(1(np.newaxis),10,2) --> (10,10,2)
tmp = x[:,np.newaxis,:] - x[np.newaxis,:,:]
print(f'Broadcasted shape:{tmp.shape}')
# 按最低维求和得到平方距离
sq_dist = np.sum((x[:,np.newaxis,:] - x[np.newaxis,:,:])**2, axis = -1)
# np.argsort可求得所有邻近（自身的最邻近即为本身）
nearest = np.argsort(sq_dist, axis=1)
# numpy的数组索引是多维数组风格
# print(nearest)
print(nearest[:,0])

# 使用np.argpartition 求得k邻近（是无序的）
knearest = np.argpartition(sq_dist, 3+1, axis=1)
# print(knearest)
print(knearest[:,:3+1])

# import seaborn as sns
import matplotlib.pyplot as plt
# 可视化k邻近
sns.set()
# plt.style.use('classic')
for i in range(x.shape[0]):
    for j in knearest[i,:3+1]:
        plt.plot(*zip(x[j], x[i]), color='black')

参考资料

[1]. Python数据科学手册,人民邮电出版社78-80