Use values in a list as index values to create new multidimensional array-CodePudding

I have an array here

A = np.array([[1,2,3],
              [4,5,6],
              [7,8,9]])

And another array that represents column index values in A

Cols,rows = np.array([[0,1],[1,2]])

I then want use those column values to index array A to end up with an array that looks like

Cols_result = [[[1,2],[4,5],[7,8]],
               [[2,3],[5,6],[8,9]]]

Row_results = [[[1,2,3],[4,5,6]],
               [[4,5,6],[7,8,9]]]

I tried using np.take() but could only get it to work properly index for row values not column values

CodePudding user response：

One way using numpy.take with numpy.swapaxes:

A = np.array([[1,2,3],[4,5,6],[7,8,9]])
ind = np.array([[0,1],[1,2]])

Output for rows:

np.take(A, ind, 0)

array([[[1, 2, 3],
        [4, 5, 6]],

       [[4, 5, 6],
        [7, 8, 9]]])

Output for columns:

np.swapaxes(np.take(A, ind, 1), 0, 1)

array([[[1, 2],
        [4, 5],
        [7, 8]],

       [[2, 3],
        [5, 6],
        [8, 9]]])