Qubix
Qubix

Reputation: 4353

Numpy apply functions to selected indices of matrix

I have two random index generators, from 0 to 3

idx_1 = np.random.choice(4, 2, replace=False)
idx_2 = np.setdiff1d(range(4), idx_1)

and I have a 3D numpy array, example_array, made of stacking 4 matrices together, and 2 functions , say func1 and func2.

I want to do something like

example_array[idx1[0]] = func1(example_array[idx1[0]])
example_array[idx1[1]] = func1(example_array[idx1[1]])
example_array[idx2[0]] = func2(example_array[idx2[0]])
example_array[idx2[1]] = func2(example_array[idx1[0]])

What is a fast way of doing that, without explicitly writing them as above?

Upvotes: 1

Views: 992

Answers (1)

Nils Werner
Nils Werner

Reputation: 36719

What about vectorizing func1() and func2() so you can do

example_array[idx1] = func1(example_array[idx1])
example_array[idx2] = func2(example_array[idx2])

Upvotes: 1

Related Questions