How to extract rows from an numpy array based on the content?
<p>As title, for example, I have an 2d numpy array, like the one below, </p>

<pre><code>[[33, 21, 1],
 [33, 21, 2],
 [32, 22, 0],
 [33, 21, 3],
 [34, 34, 1]]
</code></pre>

<p>and I want to extract these rows orderly based on the content in the first and the second column, in this case, I want to get 3 different 2d numpy arrays, as below, </p>

<pre><code>[[33, 21, 1],
 [33, 21, 2],
 [33, 21, 3]]
</code></pre>

<p>and </p>

<pre><code>[[32, 22, 0]]
</code></pre>

<p>and </p>

<pre><code>[[34, 34, 1]]
</code></pre>

<p>What function in numpy could I use to do this? I think the point is to distinguish different rows with their first and second columns. If elements in these columns are the same, then the specific rows are categorized in the same output array. <strong>I want to write a python function to do this kind of job, because I could have a much more bigger array than the one above.</strong> Feel free to give me advice, thank you.</p>

<p>You would use boolean indexing to do this. To obtain the three examples you give (in the same order as you posted them, where <code>x</code> is your original 2d array), you could write:</p>

<pre><code>numpy.atleast_2d( x[ x[:,1]==21 ] )
numpy.atleast_2d( x[ x[:,2]==0 ] )
numpy.atleast_2d( x[ x[:,2]==1 ] )
</code></pre>

<p>The first should be interpreted as saying 'extract the rows of <code>x</code> where the element in the second column equals 21' and so on. There is a page in the scipy docs that explains how to use indexing in numpy <a href="http://docs.scipy.org/doc/numpy/reference/arrays.indexing.html" rel="nofollow" title="here">here</a>. Since you required that the returned arrays all be 2D, I have used the <code>atleast_2d</code> function.</p>

python arrays numpy