python - 从行中元素满足条件的数组中获取索引

我想找到满足条件的数组的索引。

我有一个numpy.ndarray B：
（m =行数= 8，
3列）

array([[ 0.,  0.,  0.],
   [ 0.,  0.,  0.],
   [ 0.,  0.,  1.],
   [ 0.,  1.,  1.],
   [ 0.,  1.,  0.],
   [ 1.,  1.,  0.],
   [ 1.,  1.,  1.],
   [ 1.,  0.,  1.],
   [ 1.,  0.,  0.]])

对于每一列，我想找到元素满足以下条件的行的索引：
对于列中的col：
对于所有行= 1,2，..，m-1，B（row，col）= 1和B（row + 1，col）= 1；对于行= 0和m，B（row，col）= 1。

因此，理想的结果是：

Sets1=[[5, 6, 7, 8], [3, 4, 5], [2, 6]]

到目前为止，我已经尝试过了：

Sets1=[]
for j in range(3):
    Sets1.append([i for i, x in enumerate(K[1:-1]) if B[x,j]==1 and B[x+1,j]==1])

但这只是条件的第一部分，并给出以下错误输出，因为它采用了新集合的索引。因此，它实际上应为加1。

Sets1= [[4, 5, 6], [2, 3, 4], [1, 5]]

条件的第二部分也适用于索引0和m。还没有包括在内..

编辑：我通过写i + 1固定了加号1部分，并通过添加以下if语句尝试了条件的第二部分：

Sets1=[]
for j in range(3):
    Sets1.append([i+1 for i, xp in enumerate(K[1:-1]) if B[xp,j]==1 and B[xp+1,j]==1])
    if B[0,j]==1: Sets1[j].append(0)
    if B[(x-1),j]==1: Sets1[j].append(x-1)

哪个起作用，因为它给出以下输出：

Sets1= [[5, 6, 7, 8], [3, 4, 5], [2, 6]]

所以现在我只需要在条件的第一部分（在if语句之前）将+1添加到列表的元素中...

我非常感谢您的帮助！

最佳答案

您可以使用布尔掩码和np.where完成此操作

一，面膜：

c1 = (x==1)
c2 = (np.roll(x, -1, axis=0) != 0)
c3 = (x[-1] == 1)

c1 & (c2 | c3)

array([[False, False, False],
       [False, False, False],
       [False, False,  True],
       [False,  True, False],
       [False,  True, False],
       [ True,  True, False],
       [ True, False,  True],
       [ True, False, False],
       [ True, False, False]])

np.where获取索引：

>>> np.where(c1 & (c2 | c3))

(array([2, 3, 4, 5, 5, 6, 6, 7, 8], dtype=int64),
 array([2, 1, 1, 0, 1, 0, 2, 0, 0], dtype=int64))

如果您确实希望将结果作为输出中的列表：

s = np.where(c1 & (c2 | c3))
[list(s[0][s[1]==i]) for i in range(x.shape[1])]

# [[5, 6, 7, 8], [3, 4, 5], [2, 6]]

关于python - 从行中元素满足条件的数组中获取索引，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/51101791/