我有一个3D numpy数组看起来像这样
shape(3,1000,100)
[[[2,3,0,2,6,...,0,-1,-1,-1,-1,-1],
[1,4,6,1,4,5,3,...,1,2,6,-1,-1],
[7,4,6,3,1,0,1,...,2,0,8,-1,-1],
...
[8,7,6,4,...,2,4,5,2,1,-1]],
...,
[1,5,6,7,...,0,0,0,0,1]]]
阵列的每个通道以0或多个（小于70我确定） -1 结束。
目前，我想在每个泳道的 -1 之前仅选择30个值，以生成形状为(3,1000,30)的原始numpy数组的子集
应该是这样的，
[[[...,0],
[...,1,2,6],
[...,2,0,8],
...
[...,2,4,5,2,1]],
...,
[...,0,0,0,0,1]]]
是否可以使用一些numpy功能？希望没有for循环：）
答案 0 :(得分：12)
这是一个利用broadcasting和advanced-indexing -
def preceedingN(a, N):
# mask of value (minus 1 here) to be found
mask = a==-1
# Get the first index with the value along the last axis.
# In case its not found, choose the last index
idx = np.where(mask.any(-1), mask.argmax(-1), mask.shape[-1])
# Get N ranged indices along the last axis
ind = idx[...,None] + np.arange(-N,0)
# Finally advanced-index and get the ranged indexed elements as the o/p
m,n,r = a.shape
return a[np.arange(m)[:,None,None], np.arange(n)[:,None], ind]
示例运行 -
可重复输入的设置：
import numpy as np
# Setup sample input array
np.random.seed(0)
m,n,r = 2,4,10
a = np.random.randint(11,99,(m,n,r))
# Select N elements off each row
N = 3
idx = np.random.randint(N,a.shape[-1]-1,(m,n))
a[idx[...,None] < np.arange(a.shape[-1])] = -1
a[0,0] = range(r) # set first row of first 2D slice to range (no -1s there)
输入，输出：
>>> a
array([[[ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
[81, 23, 69, 76, 50, 98, 57, -1, -1, -1],
[88, 83, 20, 31, 91, 80, 90, 58, -1, -1],
[60, 40, 30, 30, 25, 50, 43, 76, -1, -1]],
[[43, 42, 85, 34, 46, 86, 66, 39, -1, -1],
[11, 47, 64, 16, -1, -1, -1, -1, -1, -1],
[42, 12, 76, 52, 68, 46, 22, 57, -1, -1],
[25, 64, 23, 53, 95, 86, 79, -1, -1, -1]]])
>>> preceedingN(a, N=3)
array([[[ 7, 8, 9],
[50, 98, 57],
[80, 90, 58],
[50, 43, 76]],
[[86, 66, 39],
[47, 64, 16],
[46, 22, 57],
[95, 86, 79]]])
答案 1 :(得分：6)
这是一个基于计算应保留的指数的想法的解决方案。我们使用numpy.argmin和numpy.nonzero来查找-1的开头或行的结尾，然后使用二维加法/减法来构建需要的30个元素的索引保持。
首先，我们创建可重现的示例数据
import numpy as np
np.random.seed(0)
a = np.random.randint(low=0, high=10, size=(3, 1000, 100))
for i in range(3):
for j in range(1000):
a[i, j, np.random.randint(low=30, high=a.shape[2]+1):] = -1
您当然可以跳过此步骤，我只是添加它以允许其他人重现此解决方案。 ：）
现在让我们逐步完成代码：
更改a的形状以简化代码
a.shape = 3 * 1000, 100
在
的每一行中查找-1的索引i = np.argmin(a, axis=1)
将a的行的任何索引替换为100
i[np.nonzero(a[np.arange(a.shape[0]), i] != -1)] = a.shape[1]
将指数翻译为1D
i = i + a.shape[1] * np.arange(a.shape[0])
计算应保留的所有元素的索引。我们将索引设为二维，以便在每个-1索引之前得到30个索引。
i = i.reshape(a.shape[0], 1) - 30 + np.arange(30).reshape(1, 30)
a.shape = 3 * 1000 * 100
执行过滤
a = a[i]
将a返回到所需的形状
a.shape = 3, 1000, 30
答案 2 :(得分：5)
这是一个使用stride_tricks来有效检索切片的方法：
import numpy as np
from numpy.lib.stride_tricks import as_strided
# create mock data
data = np.random.randint(0, 9, (3, 1000, 100))
fstmone = np.random.randint(30, 101, (3, 1000))
data[np.arange(100) >= fstmone[..., None]] = -1
# count -1s (ok, this is a bit wasteful compared to @Divakar's)
aux = np.where(data[..., 30:] == -1, 1, -100)
nmones = np.maximum(np.max(np.cumsum(aux[..., ::-1], axis=-1), axis=-1), 0)
# slice (but this I'd expect to be faster)
sliceable = as_strided(data, data.shape[:2] + (71, 30),
data.strides + data.strides[2:])
result = sliceable[np.arange(3)[:, None], np.arange(1000)[None, :], 70-nmones, :]
基准测试，最佳解决方案是@Divakar和我的混合，@Florian Rhiem也非常快：
import numpy as np
from numpy.lib.stride_tricks import as_strided
# create mock data
data = np.random.randint(0, 9, (3, 1000, 100))
fstmone = np.random.randint(30, 101, (3, 1000))
data[np.arange(100) >= fstmone[..., None]] = -1
def pp(data, N):
# count -1s
aux = np.where(data[..., N:] == -1, 1, -data.shape[-1])
nmones = np.maximum(np.max(np.cumsum(aux[..., ::-1], axis=-1), axis=-1), 0)
# slice
sliceable = as_strided(data, data.shape[:2] + (data.shape[-1]-N+1, N),
data.strides + data.strides[2:])
return sliceable[np.arange(data.shape[0])[:, None],
np.arange(data.shape[1])[None, :],
data.shape[-1]-N-nmones, :]
def Divakar(data, N):
# mask of value (minus 1 here) to be found
mask = data==-1
# Get the first index with the value along the last axis.
# In case its not found, choose the last index
idx = np.where(mask.any(-1), mask.argmax(-1), mask.shape[-1])
# Get N ranged indices along the last axis
ind = idx[...,None] + np.arange(-N,0)
# Finally advanced-index and get the ranged indexed elements as the o/p
m,n,r = data.shape
return data[np.arange(m)[:,None,None], np.arange(n)[:,None], ind]
def combined(data, N):
# mix best of Divakar's and mine
mask = data==-1
idx = np.where(mask.any(-1), mask.argmax(-1), mask.shape[-1])
sliceable = as_strided(data, data.shape[:2] + (data.shape[-1]-N+1, N),
data.strides + data.strides[2:])
return sliceable[np.arange(data.shape[0])[:, None],
np.arange(data.shape[1])[None, :],
idx-N, :]
def fr(data, N):
data = data.reshape(-1, data.shape[-1])
i = np.argmin(data, axis=1)
i[np.nonzero(data[np.arange(data.shape[0]), i] != -1)] = data.shape[1]
i = i + data.shape[1] * np.arange(data.shape[0])
i = i.reshape(data.shape[0], 1) - N + np.arange(N).reshape(1, N)
data.shape = -1,
res = data[i]
res.shape = 3, 1000, 30
return res
print(np.all(combined(data, 30) == Divakar(data, 30)))
print(np.all(combined(data, 30) == pp(data, 30)))
print(np.all(combined(data, 30) == fr(data, 30)))
from timeit import timeit
for func in pp, Divakar, combined, fr:
print(timeit('f(data, 30)', number=100, globals={'f':func, 'data':data}))
示例输出：
True
True
True
0.2767702739802189
0.13680238201050088
0.060565065999981016
0.0795100320247002