Getting median of portion of array according to bool python

Question

I have two arrays of the same length, the first one is a boolean array, the second one contains the corresponding values.

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]

I want to return an array of medians containing the median values corresponding to each portions of 1 in the boolean matrix.

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]

My unesthetic approach is to do:

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)

It works well but it's not very pythonic and very slow since I work with very big arrays.

MB-F · Accepted Answer · 2017-05-24 14:18:50Z

2

We can use np.diff to find transitions between 0 and 1. Then loop over pairs of 0/1 and 1/0 transitions and take the median from all values inbetween.

The resulting loop iterates over each group of ones.

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. ]

edited May 24, 2017 at 14:18

answered May 24, 2017 at 14:12

MB-F

23.8k5 gold badges71 silver badges127 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Divakar Over a year ago

Seems like the obvious one for this problem.

Antlafe Over a year ago

Thanks ! It's about 100time faster for 10 000 elements arrays.

Collectives™ on Stack Overflow

Getting median of portion of array according to bool python

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related