Membership checking in Numpy ndarray
<p>I have written a script that evaluates if some entry of <code>arr</code> is in <code>check_elements</code>. My approach does <em>not</em> compare single entries, but whole vectors inside of <code>arr</code>. Thus, the script checks if <code>[8, 3]</code>, <code>[4, 5]</code>, ... is in <code>check_elements</code>.</p>

<p>Here's an example: </p>

<pre><code>import numpy as np

# arr.shape -&gt; (2, 3, 2)
arr = np.array([[[8,  3],
                 [4,  5],
                 [6,  2]],

                [[9,  0],
                 [1, 10],
                 [7, 11]]])

# check_elements.shape -&gt; (3, 2)
# generally: (n, 2)
check_elements = np.array([[4, 5], [9, 0], [7, 11]])

# rslt.shape -&gt; (2, 3)
rslt = np.zeros((arr.shape[0], arr.shape[1]), dtype=np.bool)

for i, j in np.ndindex((arr.shape[0], arr.shape[1])):
    if arr[i, j] in check_elements:   # &lt;-- condition is checked against
                                      #     the whole last dimension
        rslt[i, j] = True
    else:
        rslt[i, j] = False
</code></pre>

<p>Now:</p>

<pre><code>print(rslt)
</code></pre>

<p>...would print:</p>

<pre><code>[[False  True False]
 [ True False  True]]
</code></pre>

<p>For getting the indices of I use:</p>

<pre><code>print(np.transpose(np.nonzero(rslt)))
</code></pre>

<p>...which prints the following:</p>

<pre><code>[[0 1]    # arr[0, 1] -&gt; [4, 5] -&gt; is in check_elements
 [1 0]    # arr[1, 0] -&gt; [9, 0] -&gt; is in check_elements
 [1 2]]   # arr[1, 2] -&gt; [7, 11] -&gt; is in check_elements
</code></pre>

<p>This task would be easy and performant if I would check a condition on single values, like <code>arr &gt; 3</code> or <code>np.where(...)</code>, but I am <em>not</em> interested in single values. I want to check a condition against the whole last dimension (or slices of it).</p>

<p><strong>My question is</strong>: is there a faster way to achieve the same result? Am I right that vectorized attempts and things like <code>np.where</code> can <em>not</em> be used for my problem, because they always operate on single values and not on a whole dimension or slices of that dimension?</p>

<p>The <a href="https://github.com/EelcoHoogendoorn/Numpy_arraysetops_EP" rel="nofollow">numpy_indexed</a> package (disclaimer: I am its author) contains functionality to perform these kind of queries; specifically, containment relations for nd (sub)arrays:</p>

<pre><code>import numpy_indexed as npi
flatidx = npi.indices(arr.reshape(-1, 2), check_elements)
idx = np.unravel_index(flatidx, arr.shape[:-1])
</code></pre>

<p>Note that the implementation is fully vectorized under the hood. </p>

<p>Also, note that with this approach, the order of the indices in idx match with the order of check_elements; the first item in idx are the row and col of the first item in check_elements. This information is lost when using an approach along the lines you posted above, or when using one of the alternative suggested answers, which will give you the idx sorted by their order of appearance in arr instead, which is often undesirable.</p>

python numpy indexing