根据bool python获取数组部分的中位数

时间:2017-05-24 13:48:45

标签: python numpy median

我有两个长度相同的数组,第一个是布尔数组,第二个包含相应的值。

flag   = [0,0,0,1,1,0,0,0,1,1,1,1,0,1,1]
values = [1,5,6,8,5,6,2,0,1,9,3,8,3,6,2]

我想返回一个中位数组,其中包含与布尔矩阵中每个1的部分对应的中值。

e.g。

flag   = [0,0,0,1,  1,  0,0,0, 1,  1,  1,  1, 0,1,1]
result = [0,0,0,6.5,6.5,0,0,0,5.5,5.5,5.5,5.5,0,4,4]

我的不明智的做法是:

result = np.zeros(values.shape[0])
vect = []
idx = []
for n in np.arange(result.size):
    if flag[n] > 0:
        vect.append(values[n])
        idx.append(n)
    elif flag[n] == 0:
        result[idx] = np.median(vect)
        vect = []
        idx = []
    result[idx] = np.median(vect)

它工作得很好,但由于我使用非常大的数组,所以它不是非常pythonic和非常慢。

1 个答案:

答案 0 :(得分:2)

我们可以使用np.diff来查找0和1之间的转换。然后循环遍历0/1和1/0转换对,并从中间的所有值中取出中值。

结果循环遍历每组。

flag   = [0,0,0,1,1,0,0,0,1,1,1,1,0,1,1]
values = [1,5,6,8,5,6,2,0,1,9,3,8,3,6,2]

d = np.diff(np.concatenate([[0], flag, [0]]))  # Add and append a 0 so the procedure also works if flags start or end with 1.

begin = np.flatnonzero(d==1)
end = np.flatnonzero(d==-1)

result = np.zeros_like(values, dtype=float)

for a, b in zip(begin, end):
    result[a:b] = np.median(values[a:b])

print(result)
# [ 0.   0.   0.   6.5  6.5  0.   0.   0.   5.5  5.5  5.5  5.5  0.   4.   4. ]