Indexing the last dimension of a 3D array with a 2D integer array <p>I have one 3D <code>data = NxMxD</code> numpy array, and another 2D <code>idx = NxM</code> integer array which values are in the range of <code>[0, D-1]</code>. I want to perform basic updates to each <code>data = NxM</code> entry at the depth given by the <code>idx</code> array at that position.</p> <p>For example, for <code>N = M = D = 2</code>:</p> <pre><code>data = np.zeros((2,2,2)) idx = np.array([[0,0],[1, 1]], int) </code></pre> <p>And I want to perform a simple operation like:</p> <pre><code>data[..., idx] += 1 </code></pre> <p>My expected output would be:</p> <pre><code>&gt;&gt;&gt; data array([[[ 1., 0.], [ 1., 0.]], [[ 0., 1.], [ 0., 1.]]]) </code></pre> <p><code>idx</code> indicates for each 2D coordinate which <code>D</code> should be updated. The above operation doesn't work. </p> <p>I've found <a href="https://stackoverflow.com/questions/29098417/numpy-get-2d-array-where-last-dimension-is-indexed-according-to-a-2d-array">this approach</a> in SO which solves the indexing problem by using:</p> <pre><code>data[np.arange(N)[:, None], np.arange(M)[None, :], idx] += 1 </code></pre> <p>It works fine, but looks pretty horrible needing to manually index the whole matrices for what it seems a pretty simple operation (using one matrix as a index mask for the last channel).</p> <p>Is there any better solution?</p>