Numpy中的索引方式

每次看到numpy数组和PyTorch Tensor的复杂的索引操作时就很懵，在这里整理一下numpy中的常用的索引方式(numpy中的索引方式和pytorch很相似)，以例子为主。

整数索引

import numpy as np


a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
b = a[1:3, 1:3]
c = a[1:3, [1, 2]]
d = a[[1, 2], [1, 2]]
e = a[[1, 2]]
f = a[1:3]

print('='*20, 'a', '='*20)
print(a)
print('='*20, 'b', '='*20)
print(b)
print('='*20, 'c', '='*20)
print(c)
print('='*20, 'd', '='*20)
print(d)
print('='*20, 'e', '='*20)
print(e)
print('='*20, 'f', '='*20)
print(f)

输出结果

==================== a ====================
[[1 2 3]
 [4 5 6]
 [7 8 9]]
==================== b ====================
[[5 6]
 [8 9]]
==================== c ====================
[[5 6]
 [8 9]]
==================== d ====================
[5 9]
==================== e ====================
[[4 5 6]
 [7 8 9]]
==================== f ====================
[[4 5 6]
 [7 8 9]]

bool 索引

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

print('='*20, 'a', '='*20)
print(a)
print('='*20, 'b', '='*20)
print(b)

输出结果

==================== a ====================
[[1 2 3]
 [4 5 6]
 [7 8 9]]
==================== b ====================
[6 7 8 9]

参考资料

• https://www.runoob.com/numpy/numpy-advanced-indexing.html