Match along last axis in numpy array

问题:

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