python中的数组选择

在python中，从索引可被某些数字整除的矩阵中选择行的最有效方法是什么（使用numpy、torch或any）

例1：

mat=[[1,2,3],[4,5,6],[7,8,9],[10,11,12],[13,14,15]]
num=3
output=[[1,2,3],[10,11,12]] #Rows 0,3

例2：

mat=[[1,2,3],[4,5,6],[7,8,9],[10,11,12],[13,14,15]]
num=2
output=[[1,2,3],[7,8,9],[13,14,15]] #Rows 0,2,4

---

如果您只需要一种普通的python方法，我建议您使用列表理解。使用enumerate可以查看索引是否为所需的索引之一

mat = [[1,2,3],[4,5,6],[7,8,9],[10,11,12],[13,14,15]]
num = 2

out = [row  for i, row in enumerate(mat) if not i % num]

>>> [[1, 2, 3], [7, 8, 9], [13, 14, 15]]

此外，如果您想要numpy解决方案，您可以使用：

import numpy as np
mat = np.matrix(mat)
out = mat[0::num, ]
>>> [[1, 2, 3], [7, 8, 9], [13, 14, 15]]

---

让我们做个实验吧！让

a、b、c

be

1000x1000

矩阵分别位于

NumPy

、

PyTorch

和simple

Python

列表中。 使用Intel（R）Xeon（R）CPU@2.00GHz CPU和特斯拉V100-SXM2 GPU.PyTorch版本1.6.0 w/CUDA 10.1：

import torch
import numpy as np

k=1000
a = torch.arange(k**2, device='cuda').view(k, k) #On GPU
b = np.arange(k**2).reshape(k, k)
c = [[j * k + m % k for m in range(k)] for j in range(k)]

使用

IPython

内置的magic命令

%%timeit

我们可以获得对每个三人进行切片所需的时间。在这里，我们使用100000次迭代中最差的时间来避免缓存带来的噪音。我们在

范围（1100）

中使用

num

进行检查，结果如下，图形越低越好：

import matplotlib.pyplot as plt
ts_torch = []
ts_numpy = []
ts_lists = []
strides = range(2, 100)
for num in strides:
  t_torch = %timeit -o a[::num,:]
  t_numpy = %timeit -o b[::num,:]
  t_lists = %timeit -o c[::num]
  ts_torch.append(t_torch.worst)
  ts_numpy.append(t_numpy.worst)
  ts_lists.append(t_lists.worst)

plt.rcParams['figure.figsize'] = (10,10)
plt.plot(strides, ts_torch, c='r', label='PyTorch')
plt.plot(strides, ts_numpy, c='g', label='Numpy')
plt.plot(strides, ts_lists, c='b', label='Lists')
plt.legend(loc="upper right")
plt.xlabel('stride')
plt.ylabel('time [s]')
plt.show()

获胜者是（至少对于这项非100%科学的实验）NumPy和

列表

是头对头的，

列表

在这种设置下，大步剪切比

NumPy.ndarray

稍微快一点

---

仅在

numpy

上，不能替代

mat[：：num]

。它的性能超过了列表上的

mat[：：num]

，甚至超过了数百万：）更重要的是，基础研究是无可替代的。你在学习的这一阶段提出的问题不太可能是惊天动地的新问题。

np.matrix

是老式的，

np.array

（或者至少，

np.asarray

）更好。@mathfux

np.asarray

是

np的包装器。数组

列表也可以用同样的方法编制索引。对每个索引进行测试是非常低效的。