Question

It seems I still struggle with the “in” operator in numpy. Here’s the situation:

>>> a = np.random.randint(1, 10, (2, 2, 3))
>>> a
array([[[9, 8, 8],
        [4, 9, 1]],

       [[6, 6, 3],
        [9, 3, 5]]])

I would like to get the indexes of those triplets whose second element is in (6, 8). The way I intuitively tried is:

>>> a[:, :, 1] in (6, 8)
ValueError: The truth value of an array with more than one element...

My ultimate goal would be to insert at those positions the the number preceding those multiplied by two. Using the example above, a should become:

array([[[9, 18, 8],   #8 @ pos #2 --> replaced by 9 @ pos #1 by 2
        [4, 9, 1]],

       [[6, 12, 3],   #6 @ pos #2 --> replaced by 6 @ pos #1 by 2
        [9, 3, 5]]])

Thank you in advance for your advice and time!

Answer 1

import numpy as np
a = np.array([[[9, 8, 8],
               [4, 9, 1]],

              [[6, 6, 3],
               [9, 3, 5]]])

ind=(a[:,:,1]<=8) & (a[:,:,1]>=6)
a[ind,1]=a[ind,0]*2
print(a)

yields

[[[ 9 18  8]
  [ 4  9  1]]

 [[ 6 12  3]
  [ 9  3  5]]]

---

If you wish to check for membership in a set which is not a simple range, then I like both mac’s idea of using a Python loop and bellamyj’s idea of using np.in1d. Which is faster depends on the size of check_tuple:

test.py:

import numpy as np
np.random.seed(1)

N = 10
a = np.random.randint(1, 1000, (2, 2, 3))
check_tuple = np.random.randint(1, 1000, N)

def using_in1d(a):
    idx = np.in1d(a[:,:,1], check_tuple)
    idx=idx.reshape(a[:,:,1].shape)
    a[idx,1] = a[idx,0] * 2
    return a

def using_in(a):
    idx = np.zeros(a[:,:,0].shape,dtype=bool)
    for n in check_tuple:
        idx |= a[:,:,1]==n
    a[idx,1] = a[idx,0]*2
    return a

assert np.allclose(using_in1d(a),using_in(a))    

When N = 10, using_in is slightly faster:

% python -m timeit -s'import test' 'test.using_in1d(test.a)'
10000 loops, best of 3: 156 usec per loop
% python -m timeit -s'import test' 'test.using_in(test.a)'
10000 loops, best of 3: 143 usec per loop

When N = 100, using_in1d is much faster:

% python -m timeit -s'import test' 'test.using_in1d(test.a)'
10000 loops, best of 3: 171 usec per loop
% python -m timeit -s'import test' 'test.using_in(test.a)'
1000 loops, best of 3: 1.15 msec per loop