Python - 在二进制numpy数组中找到一维重心的最佳方法

Question

假设我有以下Numpy数组，其中只有一个连续的1片段：

import numpy as np
x = np.array([0,0,0,0,1,1,1,0,0,0], dtype=1)

我希望找到1D质心的索引值，1D质心是指1元素的质心。可以输入以下代码：

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

(当1的数量为偶数时，这会导致粗略的近似。) 有没有更好的方法来完成这个任务，比如一个计算效率更高的一行代码？

- fmonegaglia

Answer 1

你不能简单地执行以下操作吗？

center_of_mass = (x*np.arange(len(x))).sum()/x.sum() # 5

%timeit center_of_mass = (x*arange(len(x))).sum()/x.sum()
# 100000 loops, best of 3: 10.4 µs per loop

- Julien Spronck

Answer 2

作为一种方法，我们可以获取非零索引，并将它们的平均值作为质心，如下所示 -

np.flatnonzero(x).mean()

这里是另一种方法，使用移位数组比较来获取该切片的起始和结束索引，并获取这些索引的平均值以确定重心，代码如下 -

np.flatnonzero(x[:-1] != x[1:]).mean()+0.5

运行时测试 -

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

我们可以使用np.nonzero()[0]来替换np.flatnonzero，并且使用np.sum代替np.mean，从而在这里提高性能。

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

另外一种方法是，我们可以存储起始和结束索引，然后将它们相加以获得质心。这样做更加高效，因为我们避免了调用np.mean函数，代码如下 -

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

时间 -

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

在这里，我们可以尝试使用np.nonzero()[0]进行实验。

- Divakar