Find index of multiple queries in a multidimensional numpy array
<p>I am looking for a way to find indices of an array of queries in a multidimensional array. For example:</p>

<pre><code>arr = np.array([[17,  5, 19,  9],
   [18, 13,  3,  7],
   [ 8,  1,  4,  2],
   [ 8,  9,  7, 19],
   [ 6, 11,  8,  5],
   [11, 16, 13, 18],
   [ 0,  1,  2,  9],
   [ 1,  7,  4,  6]])
</code></pre>

<p>I can find indices for one query:</p>

<pre><code>np.where(arr==1)
# (array([2, 6, 7]), array([1, 1, 0]))
</code></pre>

<p>Is there any numpy solution to do this for multiple values replacing the following <code>for</code> loop?</p>

<pre><code>for q in queries:
    np.where(arr==q)
</code></pre>

<p>If both the array and queries were one-dimensional, I could use <code>searchsorted</code> as <a href="http://stackoverflow.com/a/9566681/3349443">this answer</a> but it doesn't work for multidimensional arrays.</p>

<p>You can get the indices of each matching value by zipping the results of your <code>where</code> function and then using the <code>*</code> dereferencing operator.</p>

<pre><code>arr = np.array([[17,  5, 19,  9],
                [18, 13,  3,  7],
                [ 8,  1,  4,  2],  # (2, 1)
                [ 8,  9,  7, 19],
                [ 6, 11,  8,  5],
                [11, 16, 13, 18],
                [ 0,  1,  2,  9],  # (6, 1)
                [ 1,  7,  4,  6]])  # (7, 0)

&gt;&gt;&gt; zip(*np.where(arr == 1))
[(2, 1), (6, 1), (7, 0)]
</code></pre>

<p>I'm not sure what your intended output is, but you can use a dictionary comprehension to show the index location for a given set of numbers, e.g.:</p>

<pre><code>&gt;&gt;&gt; {n: zip(*np.where(arr == n)) for n in range(5)}
{0: [(6, 0)],
 1: [(2, 1), (6, 1), (7, 0)],
 2: [(2, 3), (6, 2)],
 3: [(1, 2)],
 4: [(2, 2), (7, 2)]}
</code></pre>

python arrays numpy multidimensional-array