извлечения индексов из многомерной матрицы с использованием условия, max

Question

У меня есть 5-мерная матрица.

Я хочу иметь возможность извлекать индексы первого значения в каждой строке определенного (2d) среза, который удовлетворяет условию, а затем использовать эти индексы для извлечения значения соответствующих индексов в другом фрагменте.

Вот мой пример:

In [3]: g = np.arange(48400).reshape(20,11,11,2,10) 

два ломтика я работаю с являются:

In [4]: sliceA = g[0,:,:,0,0] 

In [5]: sliceA 
Out[5]: 
array([[ 0, 20, 40, 60, 80, 100, 120, 140, 160, 180, 200], 
     [ 220, 240, 260, 280, 300, 320, 340, 360, 380, 400, 420], 
     [ 440, 460, 480, 500, 520, 540, 560, 580, 600, 620, 640], 
     [ 660, 680, 700, 720, 740, 760, 780, 800, 820, 840, 860], 
     [ 880, 900, 920, 940, 960, 980, 1000, 1020, 1040, 1060, 1080], 
     [1100, 1120, 1140, 1160, 1180, 1200, 1220, 1240, 1260, 1280, 1300], 
     [1320, 1340, 1360, 1380, 1400, 1420, 1440, 1460, 1480, 1500, 1520], 
     [1540, 1560, 1580, 1600, 1620, 1640, 1660, 1680, 1700, 1720, 1740], 
     [1760, 1780, 1800, 1820, 1840, 1860, 1880, 1900, 1920, 1940, 1960], 
     [1980, 2000, 2020, 2040, 2060, 2080, 2100, 2120, 2140, 2160, 2180], 
     [2200, 2220, 2240, 2260, 2280, 2300, 2320, 2340, 2360, 2380, 2400]]) 

и один я сделал отдельно, затем добавляют (для иллюстрации):

In [6]: sliceB = np.array([[ 3, 12, 21, 31, 41, 51, 69, 77, 83, 91, 100], 
    ...:     [ 6, 12, 23, 33, 43, 51, 69, 77, 83, 91, 100], 
    ...:     [ 8, 12, 27, 37, 47, 51, 69, 77, 83, 91, 100], 
    ...:     [ 4, 12, 28, 38, 48, 51, 69, 77, 83, 91, 100], 
    ...:     [ 7, 12, 29, 39, 49, 51, 69, 77, 83, 91, 100], 
    ...:     [ 9, 12, 22, 32, 42, 51, 69, 77, 83, 91, 100], 
    ...:     [ 6, 12, 21, 31, 41, 51, 69, 77, 83, 91, 100], 
    ...:     [ 8, 12, 25, 35, 45, 51, 69, 77, 83, 91, 100], 
    ...:     [ 5, 12, 26, 36, 46, 51, 69, 77, 83, 91, 100], 
    ...:     [ 7, 12, 22, 32, 42, 51, 69, 77, 83, 91, 100], 
    ...:     [ 3, 12, 24, 34, 44, 51, 69, 77, 83, 91, 100]]) 

In [11]: g[0,:,:,0,1] = sliceB 

In [12]: g[0,:,:,0,1] 
Out[12]: 
array([[ 3, 12, 21, 31, 41, 51, 69, 77, 83, 91, 100], 
     [ 6, 12, 23, 33, 43, 51, 69, 77, 83, 91, 100], 
     [ 8, 12, 27, 37, 47, 51, 69, 77, 83, 91, 100], 
     [ 4, 12, 28, 38, 48, 51, 69, 77, 83, 91, 100], 
     [ 7, 12, 29, 39, 49, 51, 69, 77, 83, 91, 100], 
     [ 9, 12, 22, 32, 42, 51, 69, 77, 83, 91, 100], 
     [ 6, 12, 21, 31, 41, 51, 69, 77, 83, 91, 100], 
     [ 8, 12, 25, 35, 45, 51, 69, 77, 83, 91, 100], 
     [ 5, 12, 26, 36, 46, 51, 69, 77, 83, 91, 100], 
     [ 7, 12, 22, 32, 42, 51, 69, 77, 83, 91, 100], 
     [ 3, 12, 24, 34, 44, 51, 69, 77, 83, 91, 100]]) 

Теперь я хочу создать массив индексов первого элемента в каждой строке sliceB, который удовлетворяет условию ион (например, > = 35), то есть эти значения:

array([[ 3, 12, 21, 31, *41*, 51, 69, 77, 83, 91, 100], 
     [ 6, 12, 23, 33, *43*, 51, 69, 77, 83, 91, 100], 
     [ 8, 12, 27, *37*, 47, 51, 69, 77, 83, 91, 100], 
     [ 4, 12, 28, *38*, 48, 51, 69, 77, 83, 91, 100], 
     [ 7, 12, 29, *39*, 49, 51, 69, 77, 83, 91, 100], 
     [ 9, 12, 22, 32, *42*, 51, 69, 77, 83, 91, 100], 
     [ 6, 12, 21, 31, *41*, 51, 69, 77, 83, 91, 100], 
     [ 8, 12, 25, *35*, 45, 51, 69, 77, 83, 91, 100], 
     [ 5, 12, 26, *36*, 46, 51, 69, 77, 83, 91, 100], 
     [ 7, 12, 22, 32, *42*, 51, 69, 77, 83, 91, 100], 
     [ 3, 12, 24, 34, *44*, 51, 69, 77, 83, 91, 100]]) 

затем использовать, чтобы создать массив значений в sliceA с соответствующими индексами, то есть:

array([[ 0, 20, 40, 60, *80*, 100, 120, 140, 160, 180, 200], 
     [ 220, 240, 260, 280, *300*, 320, 340, 360, 380, 400, 420], 
     [ 440, 460, 480, *500*, 520, 540, 560, 580, 600, 620, 640], 
     [ 660, 680, 700, *720*, 740, 760, 780, 800, 820, 840, 860], 
     [ 880, 900, 920, *940*, 960, 980, 1000, 1020, 1040, 1060, 1080], 
     [1100, 1120, 1140, 1160, *1180*, 1200, 1220, 1240, 1260, 1280, 1300], 
     [1320, 1340, 1360, 1380, *1400*, 1420, 1440, 1460, 1480, 1500, 1520], 
     [1540, 1560, 1580, *1600*, 1620, 1640, 1660, 1680, 1700, 1720, 1740], 
     [1760, 1780, 1800, *1820*, 1840, 1860, 1880, 1900, 1920, 1940, 1960], 
     [1980, 2000, 2020, 2040, *2060*, 2080, 2100, 2120, 2140, 2160, 2180], 
     [2200, 2220, 2240, 2260, *2280*, 2300, 2320, 2340, 2360, 2380, 2400]]) 

Я пытался в течение нескольких часов, используя следующие функции : np.amax, np.argmax, np.where, x [x> 34] .min()

, но не может найти недостающую ссылку или комбинацию.

Я хотел бы сделать это без петель в интересах скорости.

Answer 1

я не могу проверить это прямо сейчас, но это должно быть довольно просто:

idx = np.argmax(sliceB >= 35, axis=1) # index of first occurrence of condition 
sliceA[np.arange(sliceA.shape[0]), idx] 

Answer 2

Что-то, как это должно работать:

#First sort sliceA 
tmp = np.argsort(sliceA,axis=1)   
#Mask all indices that you dont want with values larger then any in the array 
tmp[ sliceB<=34 ] = tmp.shape[-1]*2   
#Find the minimum positions 
min_pos = tmp.argmin(axis=1) 

#Finally take the slice 
print sliceA[np.arange(sliceA.shape[0]),min_pos] 
[ 80 300 500 720 940 1180 1400 1600 1820 2060 2280]