python - 在多种条件下从数组中选择行

Question

我正在尝试从大型 numpy 数组中提取行。数组的列是obs编号、组id(j)、时间id(t)和一些数据x_jt。

这是一个例子:

import numpy as np
N = 100
T = 100
X = np.vstack((np.array(range(1,N*T+1)),np.repeat(np.array(range(1,N+1)),T), np.tile(np.array(range(1,T+1)),N), np.random.randint(100,size=N*T))).T

如果我想从 X 中提取组 id = 2 的所有行，我会这样做

X[np.where(X[:,1] == 2)]

如果我想要 j = 2 或 3 的所有行，我可以扩展该代码。然而，就我而言，我有许多组 id (j) 需要提取。具体来说，我想提取 j 来自的所有行

samples = np.random.randint(N, size=N) + 1

例如，假设大小 = 5 而不是 N，并且样本 = (2,4,5,4,7)。我所追求的是遍历 X 并选择 j = 2、j = 4、j = 5、j = 4、最后 j = 7 的所有行的代码，并使用结果创建一个新数组。基本上是这样的:

result = []
for j in samples:
    result.extend(X[np.where(X[:,1] == j)])

但是，当 N 较大时，此代码会很慢。您有什么加快速度的建议吗？谢谢!

Answer 1

无需更换

这可以通过向量化函数来完成:

def contains(X, samples):
    return numpy.vectorize(lambda x: x in samples)(X)

result = X[contains(X[:, 1], set(samples)), :]

更换

如果您想通过替换来执行此操作，只需检查每个样本的一个计数，直到没有更多样本(假设顺序无关紧要)。这样您至少可以减少需要迭代矩阵的次数。

result = []
sample_counts = collections.Counter(samples)
while sum(sample_counts.itervalues()):
    # pick up one of each of the remaining samples and chain their rows
    # together in result
    s = set(key for key, value in sample_counts.iteritems() if value)
    result = itertools.chain(result, X[contains(X[:, 1], s), :])
    sample_counts -= collections.Counter(dict.fromkeys(s, 1))
# create a matrix of the final result
result = numpy.array(list(result))

在这种情况下，我能想到的唯一可能加快您已经在做的事情的方法是预分配矩阵。所以你会这样做: