extract indices from multi dimensional array using condition, max

Question

I have a 5 dimensional matrix.

I want to be able to extract the indices of the first value in each row of a certain (2d) slice that meets a condition, then use those indices to extract the value of the corresponding indices in another slice.

Here is my example:

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

the two slices I'm working with are:

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]])

and one I made separately, then added in (for illustration purposes):

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]])

Now, I want to create an array of indices of the first element in each row of sliceB that meets a condition (eg. >=35), ie these values:

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]])

then use that to create an array of values in sliceA with the corresponding indices, ie:

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]])

I have tried for hours using the following functions: np.amax, np.argmax, np.where, x[x>34].min()

but can't seem to find the missing link or combination.

I'd like to do this without loops in the interest of speed.

check `np.argwhere` and http://stackoverflow.com/questions/5957470/matlab-style-find-function-in-python/5957742#5957742 — joaquin, Sep 14 '13 at 07:37
Yes, but when doing this for each row of the array and trying to create another array in the process (because I believe it creates a reference and not another memory location, therefore is quicker), It becomes confusing. The other difference is that in the mentioned post, the questioner is looking for each value of the array that meets the condition. I'm looking for only the first (minimum) — Sammy, Sep 14 '13 at 07:54

score 3 · Accepted Answer · answered Sep 14 '13 at 14:16

3

I can't test it right now, but it should be pretty simple:

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

answered Sep 14 '13 at 14:16

Jaime

65,696
17
124
159

score 2 · Answer 2 · answered Sep 14 '13 at 13:26

Something like this should work:

#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]

extract indices from multi dimensional array using condition, max

2 Answers2