Выберите значения из набора массивов в соответствии с массивом перестановок

Question

У меня есть 3 массива numpy формы 2xN (с большим N, несколько миллионов), назовите их a1, a2, a3. Тогда у меня есть еще один массив формы Nx3, значения строк которого относятся к одному из массивов a1, a2, a3, называют его перестановками. Этот перестановками массив выглядит следующим образом: [[0, 1, 2], [1,2,0], [1,0,2], ... до N строк]

Я хочу создайте еще 3 массива numpy b1, b2, b3 формы 2xN, которые имеют содержимое исходных a1, a2, a3, но их столбцы были перегруппированы в соответствии с строками массива перестановок.

Я пробовал причудливую индексацию, укладывая 3 массива, и numpy.choose, но я не могу заставить ее работать. Я ищу решение без петель python. Любая помощь будет принята с благодарностью!

EDIT

Просто для уточнения я показываю реализацию Питона цикла, что я пытаюсь сделать:

aa = np.dstack((a1, a2, a3)) 
bb = np.empty_like(aa) 
for i, o in enumerate(permutations): 
    bb[:,i, np.arange(3)] = aa[:, i, o] 

Тогда я бы восстановить b1, b2, b3 от бб.

Answer 1

С fancy-indexing, вы могли бы сделать -

bb = aa[:,np.arange(N),permutations.T] 

Пожалуйста, обратите внимание, что это будет иметь форму (2,3,N). Таким образом, чтобы выбрать b1, b2, b3, вы могли бы сделать:

b1,b2,b3 = bb[:,0,:], bb[:,1,:], bb[:,2,:] 

Или, если вы настаиваете bb быть той же форма, как и с публикуемым кодом, вы могли бы добавить:

bb = bb.swapaxes(1,2) 

---

Вот еще один подход с использованием линейной индексации, нарезка и конечно NumPy broadcasting -

idx = permutations + 3*np.arange(N)[:,None]  
bb = aa.reshape(2,-1)[:,idx].reshape(2,N,3) 

Это создало бы bb той же формы, что и для отправленного кода цикла.

---

тест Продолжительность

In [189]: def original_app(aa,permutations): 
    ...:  bb = np.empty_like(aa) 
    ...:  for i, o in enumerate(permutations): 
    ...:   bb[:,i, np.arange(3)] = aa[:, i, o] 
    ...:  return bb 
    ...: 
    ...: 
    ...: def linear_index_app(aa,permutations): 
    ...:  idx = permutations + 3*np.arange(N)[:,None]  
    ...:  return aa.reshape(2,-1)[:,idx].reshape(2,N,3) 
    ...: 

In [190]: # Setup input arrays 
    ...: N = 10000 
    ...: a1 = np.random.rand(2,N) 
    ...: a2 = np.random.rand(2,N) 
    ...: a3 = np.random.rand(2,N) 
    ...: 
    ...: permutations = np.random.randint(0,3,(N,3)) 
    ...: aa = np.dstack((a1, a2, a3)) 


In [191]: %timeit original_app(aa,permutations) 
10 loops, best of 3: 128 ms per loop 

In [192]: %timeit aa[:,np.arange(N),permutations.T] 
1000 loops, best of 3: 972 µs per loop 

In [193]: %timeit linear_index_app(aa,permutations) 
1000 loops, best of 3: 1.02 ms per loop 

Таким образом, кажется, что fancy-indexing является лучшим из партии!