Pyhon - Best way to find the 1d center of mass in a binary numpy array
<p>Suppose I have the following Numpy array, in which I have one and only one continuous slice of <code>1</code>s:</p>

<pre><code>import numpy as np
x = np.array([0,0,0,0,1,1,1,0,0,0], dtype=1)
</code></pre>

<p>and I want to find the index of the 1D center of mass of the <code>1</code> elements. I could type the following:</p>

<pre><code>idx = np.where( x )[0]
idx_center_of_mass = int(0.5*(idx.max() + idx.min()))
# this would give 5
</code></pre>

<p>(Of course this would lead to rough approximation when the number of elements of the <code>1</code>s slice is even.)
Is there any better way to do this, like a computationally more efficient oneliner?</p>

<p>As one approach we can get the non-zero indices and get the mean of those as the center of mass, like so -</p>

<pre><code>np.flatnonzero(x).mean()
</code></pre>

<p>Here's another approach using shifted array comparison to get the start and stop indices of that slice and getting the mean of those indices for determining the center of mass, like so -</p>

<pre><code>np.flatnonzero(x[:-1] != x[1:]).mean()+0.5
</code></pre>

<p>Runtime test -</p>

<pre><code>In [72]: x = np.zeros(10000,dtype=int)

In [73]: x[100:2000] = 1

In [74]: %timeit np.flatnonzero(x).mean()
10000 loops, best of 3: 115 Âµs per loop

In [75]: %timeit np.flatnonzero(x[:-1] != x[1:]).mean()+0.5
10000 loops, best of 3: 38.7 Âµs per loop
</code></pre>

<p>We can improve the performance by some margin here with the use of <code>np.nonzero()[0]</code> to replace <code>np.flatnonzero</code> and <code>np.sum</code> in place of <code>np.mean</code> -</p>

<pre><code>In [107]: %timeit (np.nonzero(x[:-1] != x[1:])[0].sum()+1)/2.0
10000 loops, best of 3: 30.6 Âµs per loop
</code></pre>

<p>Alternatively, for the second approach, we can store the start and stop indices and then simply add them to get the center of mass for a bit more efficient approach as we would avoid the function call to <code>np.mean</code>, like so -</p>

<pre><code>start,stop = np.flatnonzero(x[:-1] != x[1:])
out = (stop + start + 1)/2.0
</code></pre>

<p>Timings -</p>

<pre><code>In [90]: %timeit start,stop = np.flatnonzero(x[:-1] != x[1:])
10000 loops, best of 3: 21.3 Âµs per loop

In [91]: %timeit (stop + start + 1)/2.0
100000 loops, best of 3: 4.45 Âµs per loop
</code></pre>

<p>Again, we can experiment with <code>np.nonzero()[0]</code> here.</p>

python arrays numpy binary boolean