Dynamically obtain a vector along 1 dimension in 'n' dimensional numpy array

Question

I have a numpy variable that can be 'n' dimensions, for example: game_board = np.zeros((4,3,3), dtype=np.int8)

I want to obtain a vector along the first dimension based on a vector choose_vector choose_vector = np.array([x,y],dtype=np.int8)

I know how i can do this statically:

game_board[:, x, y] 
# will return [0,0,0,0], the (x,y)th element from 1st dimension

but everything I have tried so far doing this using the choose_vector has not worked:

game_board[:, choose_vector] 
# returns 
[[[0 0 0]
  [0 0 0]]

 [[0 0 0]
  [0 0 0]]

 [[0 0 0]
  [0 0 0]]

 [[0 0 0]
  [0 0 0]]]

print(game_board[choose_vector])
# returns
[[0,0,0]]

how do i construct the index for game_board given choose_vector in order to get the same result as game_board[:, x, y]

I'd then like to expand it to any dimensional game board, but I can probably work it out if i know how to do the above :)

Answer 1

This might not be the cleanest solution, but it doing what you want:

import numpy as np
x,y = 0,0

game_board = np.zeros((4,3,3), dtype=np.int8)
choose_vector = np.array([x, y], dtype=np.uint8)
game_board[[np.newaxis] + choose_vector.tolist()]

The trick is, that you can "replace" the : in your static approach with a np.newaxis inside of a python list .

Answer 2

I have worked this out with help from FlashTek. Instead of using np.newaxis, using slice(None) seems to be a drop in replacement for :

import numpy as np
x,y = 0,0

game_board = np.zeros((4,3,3), dtype=np.int8)
choose_vector = np.array([x, y], dtype=np.uint8)
game_board[[slice(None)] + choose_vector.tolist()]