Numpy 多维情况下逐行挑选元素并修改

举例：

example = np.array([ [0,1,2], [3,4,5], [6,7,8], [9, 10, 11] ] )

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

 

-> output == array([1,5,6,9])

 

实现方法：

# Extract data

# Output == array([1,5,6,9])

example[np.arange(len(example)), indx] 

 

# Modify data

example[np.arange(len(example)), indx]  = np.array([8,8,8,8])

# Example after modification -> array([ [0,8,2], [3,4,8], [8,7,8], [8, 10, 11] ] )