Match along last axis in numpy array

馋奶兔 提交于 2020-05-09 17:13:19

问题


I saw a lot of articles and answers to other questions about slicing 3D lists in python, but I can't apply those methods to my case.

I have a 3D list:

list = [
    [[0, 56, 78], [4, 86, 90], [7, 87, 34]],
    [[1, 49, 76], [0, 76, 78], [8, 60, 7]], 
    [[9, 6, 58], [6, 57, 78], [10, 46, 2]]
    ]

The the last 2 values of the 3rd dimension stay constant but change every time I rerun the code. What the code needs to do is find 2 specific pairs of those last 2 values and slice from one pair to the other. So for example:

pair1 = 86, 90
pair2 = 76, 78

The output should be:

[4, 86, 90], [7, 87, 34], [1, 49, 76], [0, 76, 78]

I know how to find the 2 pairs, I'm just not sure how to slice the list. Thanks in advance for your help and leave a comment if something is unclear.


回答1:


Using numpy, you can slice the array to keep the last two elements on the last axis, find the indices where each pair takes place, flatten the result and use it to slice the array:

a = np.array(my_list) # don't call your list "list"

a_sliced = a[...,1:]

ix1 = np.flatnonzero((a_sliced  == pair1).all(-1).ravel()).item()

ix2 = np.flatnonzero((a_sliced  == pair2).all(-1).ravel()).item()

np.concatenate(a)[ix1:ix2+1]

array([[ 4, 86, 90],
       [ 7, 87, 34],
       [ 1, 49, 76],
       [ 0, 76, 78]])

a being defined as a numpy array, and both pairs defined as tuples:

pair1 = 86, 90
pair2 = 76, 78


来源:https://stackoverflow.com/questions/61291995/match-along-last-axis-in-numpy-array

易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!