Question

我正在寻找一个优雅的解决方案：

data = np.loadtxt(file)
# data[:,0] is a time
# data[:,1] is what I want to extract
mean = 0.0
count = 0
for n in xrange(np.size(data[:,0])):
    if data[n,0] >= tstart and data[n,0] <= tend:
        mean = mean + data[n,1]
        count = count + 1

mean = mean / float(count)

我猜我可以首先提取我的2D数组然后在其上应用np.mean但我觉得可能有一些列表理解的优点使这更优雅（我来自FORTRAN背景.. ）。我在想类似的东西（显然是错误的，因为i不会是索引）：

np.mean([x for x in data[i,1] for i in data[:,0] if i >= tstart and i <= tend])

Answer 1

在numpy而不是listcomp中，您可以使用列表和数组进行索引。具体来说，假设我们有一个像您正在使用的2D阵列：

>>> import numpy as np
>>> data = np.arange(20).reshape(10, 2)
>>> data
array([[ 0,  1],
       [ 2,  3],
       [ 4,  5],
       [ 6,  7],
       [ 8,  9],
       [10, 11],
       [12, 13],
       [14, 15],
       [16, 17],
       [18, 19]])

我们可以得到第一列：

>>> ts = data[:,0] 
>>> ts
array([ 0,  2,  4,  6,  8, 10, 12, 14, 16, 18])

创建一个与我们想要的术语相对应的布尔数组：

>>> (ts >= 2) & (ts <= 6)
array([False,  True,  True,  True, False, False, False, False, False, False], dtype=bool)

然后我们可以使用它来选择我们感兴趣的列的元素：

>>> data[:,1][(ts >= 2) & (ts <= 6)]
array([3, 5, 7])

最后采取其意思：

>>> np.mean(data[:,1][(ts >= 2) & (ts <= 6)])
5.0

或者，在一行中：

>>> np.mean(data[:,1][(data[:,0] >= 2) & (data[:,0] <= 6)])
5.0

[编辑：data[:,1][(data[:,0] >= 2) & (data[:,0] <= 6)].mean()也会奏效;我总是忘记你可以使用方法。]

优雅的列表理解，根据另一个维度中的值提取数组一维中的值

1 个答案: