Identification of rows containing column median in numpy matrix of cum percentiles
<p>Consider the matrix <code>quantiles</code> that's a subset <code>[:8,:3,0]</code> of a 3D matrix with shape <code>(10,355,8)</code>. </p>

<pre><code>quantiles = np.array([
       [ 1.        ,  1.        ,  1.        ],
       [ 0.63763978,  0.61848863,  0.75348137],
       [ 0.43439645,  0.42485407,  0.5341457 ],
       [ 0.22682343,  0.18878366,  0.25253915],
       [ 0.16229408,  0.12541476,  0.15263742],
       [ 0.12306046,  0.10372971,  0.09832783],
       [ 0.09271845,  0.08209844,  0.05982584],
       [ 0.06363636,  0.05471266,  0.03855727]])
</code></pre>

<p>I want a boolean output of the same shape as the <code>quantiles</code> matrix where <code>True</code> marks the row in which the median is located:</p>

<pre><code>In [21]: medians
Out[21]:
array([[False, False, False],
       [ True, True,  False],
       [False,  False, True],
       [False, False, False],
       [False, False, False],
       [False, False, False],
       [False, False, False],
       [False, False, False]], dtype=bool)
</code></pre>

<p>To achieve this, I have the following algorithm in mind:</p>

<p>1) Identify the entries that are greater than <code>.5</code>:</p>

<pre><code>In [22]: quantiles&gt;.5
Out[22]:
array([[ True,  True,  True],
       [ True,  True,  True],
       [False, False,  True],
       [False, False, False],
       [False, False, False],
       [False, False, False],
       [False, False, False],
       [False, False, False]], dtype=bool)
</code></pre>

<p>2) Considering only the values subset by the <code>quantiles&gt;.5</code> operation, mark the row that minimizes the <code>np.abs</code> distance between the entry and <code>.5</code>. Torturing the terminology a bit, I wish to intersect the two matrices of <code>np.argmin(np.abs(quantiles-.5),axis=0)</code> and <code>quantiles&gt;.5</code> to get the above result. However, I cannot for my life figure out a way to perform the <code>np.argmin</code> on the subset and retain the shape of the <code>quantile</code> matrix. </p>

<p>PS. Yes, there is a similar question <a href="http://stackoverflow.com/questions/33658580/find-the-median-of-each-row-of-a-2-dimensional-array-in-python">here</a> but it doesn't implement my algorithm which could be, I think, more efficient on a larger scale</p>

<h2>Approach #1</h2>

<p>Here's an approach using <a href="http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html" rel="nofollow"><code>broadcasting</code></a> and some masking trick -</p>

<pre><code># Mask of quantiles lesser than or equal to 0.5 to select the invalid ones
mask1 = quantiles&lt;=0.5

# Since we are dealing with quantiles, the elems won't be &gt; 1, 
# which can be leveraged here as we will add 1s to invalid elems, and 
# then look for argmin across each col
min_idx = (np.abs(quantiles-0.5)+mask1).argmin(0)

# Let some broadcasting magic happen here!
out = min_idx == np.arange(quantiles.shape[0])[:,None]
</code></pre>

<p><strong>Step-by-step run</strong></p>

<p>1) Input :</p>

<pre><code>In [37]: quantiles
Out[37]: 
array([[ 1.        ,  1.        ,  1.        ],
       [ 0.63763978,  0.61848863,  0.75348137],
       [ 0.43439645,  0.42485407,  0.5341457 ],
       [ 0.22682343,  0.18878366,  0.25253915],
       [ 0.16229408,  0.12541476,  0.15263742],
       [ 0.12306046,  0.10372971,  0.09832783],
       [ 0.09271845,  0.08209844,  0.05982584],
       [ 0.06363636,  0.05471266,  0.03855727]])
</code></pre>

<p>2) Run the code :</p>

<pre><code>In [38]: mask1 = quantiles&lt;=0.5
    ...: min_idx = (np.abs(quantiles-0.5)+mask1).argmin(0)
    ...: out = min_idx == np.arange(quantiles.shape[0])[:,None]
    ...: 
</code></pre>

<p>3) Analyze output at each step :</p>

<pre><code>In [39]: mask1
Out[39]: 
array([[False, False, False],
       [False, False, False],
       [ True,  True, False],
       [ True,  True,  True],
       [ True,  True,  True],
       [ True,  True,  True],
       [ True,  True,  True],
       [ True,  True,  True]], dtype=bool)

In [40]: np.abs(quantiles-0.5)+mask1
Out[40]: 
array([[ 0.5       ,  0.5       ,  0.5       ],
       [ 0.13763978,  0.11848863,  0.25348137],
       [ 1.06560355,  1.07514593,  0.0341457 ],
       [ 1.27317657,  1.31121634,  1.24746085],
       [ 1.33770592,  1.37458524,  1.34736258],
       [ 1.37693954,  1.39627029,  1.40167217],
       [ 1.40728155,  1.41790156,  1.44017416],
       [ 1.43636364,  1.44528734,  1.46144273]])

In [41]: (np.abs(quantiles-0.5)+mask1).argmin(0)
Out[41]: array([1, 1, 2])

In [42]: min_idx == np.arange(quantiles.shape[0])[:,None]
Out[42]: 
array([[False, False, False],
       [ True,  True, False],
       [False, False,  True],
       [False, False, False],
       [False, False, False],
       [False, False, False],
       [False, False, False],
       [False, False, False]], dtype=bool)
</code></pre>

<p><strong>Performance boost</strong> : Following the comments, it seems to get min_idx, we can just do :</p>

<pre><code>min_idx = (quantiles+mask1).argmin(0)
</code></pre>

<h2>Approach #2</h2>

<p>This is focused on memory efficiency.</p>

<pre><code># Mask of quantiles greater than 0.5 to select the valid ones
mask = quantiles&gt;0.5

# Select valid elems
vals = quantiles.T[mask.T]

# Get vald count per col
count = mask.sum(0)

# Get the min val per col given the mask
minval = np.minimum.reduceat(vals,np.append(0,count[:-1].cumsum()))

# Get final boolean array by just comparing the min vals across each col
out = np.isclose(quantiles,minval)
</code></pre>

python arrays numpy boolean median