Count values in overlapping sliding windows in python

Viewed 1269

Given an array, a, of sorted values, and an array of ranges, bins, what is the most efficient way to count how many values in a fall within each range, rng, in bins?

Currently I am doing the following:

def sliding_count(a, end, window, start=0, step=1):
    bins = [(x, x + window) for x in range(start, (end + 1) - window, step)]
    counts = np.zeros(len(bins))
    for i, rng in enumerate(bins):
        count = len(a[np.where(np.logical_and(a>=rng[0], a<=rng[1]))])
        counts[i] = count
    return counts

a = np.array([1, 5, 8, 11, 14, 19])
end = 20
window = 10
sliding_count(a, end, window)

Which returns the expected array

array([3., 4., 3., 3., 4., 4., 3., 3., 3., 3., 3.])

But I feel like there must be a more effective way of doing this?

3 Answers
import numpy as np

def alt(a, end, window, start=0, step=1):
    bin_starts = np.arange(start, end+1-window, step)
    bin_ends = bin_starts + window
    last_index = np.searchsorted(a, bin_ends, side='right')
    first_index = np.searchsorted(a, bin_starts, side='left')
    return  last_index - first_index

def sliding_count(a, end, window, start=0, step=1):
    bins = [(x, x + window) for x in range(start, (end + 1) - window, step)]
    counts = np.zeros(len(bins))
    for i, rng in enumerate(bins):
        count = len(a[np.where(np.logical_and(a>=rng[0], a<=rng[1]))])
        counts[i] = count
    return counts

a = np.array([1, 5, 8, 11, 14, 19])
end = 20
window = 10

print(sliding_count(a, end, window))
# [3. 4. 3. 3. 4. 4. 3. 3. 3. 3. 3.]

print(alt(a, end, window))
# [3 4 3 3 4 4 3 3 3 3 3]

How alt works:

Generate the starting and ending values of the bins:

In [73]: bin_starts = np.arange(start, end+1-window, step); bin_starts
Out[73]: array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10])

In [74]: bin_ends = bin_starts + window; bin_ends
Out[74]: array([10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20])

Since a is in sorted order, you can use np.searchsorted to find the first and last index in bin_starts and bin_ends where each value in a fits:

In [75]: last_index = np.searchsorted(a, bin_ends, side='right'); last_index
Out[75]: array([3, 4, 4, 4, 5, 5, 5, 5, 5, 6, 6])

In [76]: first_index = np.searchsorted(a, bin_starts, side='left'); first_index
Out[76]: array([0, 0, 1, 1, 1, 1, 2, 2, 2, 3, 3])

The count is simply the difference in indices:

In [77]: last_index - first_index
Out[77]: array([3, 4, 3, 3, 4, 4, 3, 3, 3, 3, 3])

Here is a perfplot comparing the performance of alt versus sliding_count as a function of the length of a:

import perfplot

def make_array(N):
    a = np.random.randint(10, size=N)
    a = a.cumsum()
    return a

def using_sliding(a):
    return sliding_count(a, end, window)

def using_alt(a):
    return alt(a, end, window)

perfplot.show(
    setup=make_array,
    kernels=[using_sliding, using_alt],
    n_range=[2**k for k in range(22)],
    logx=True,
    logy=True,
    xlabel='len(a)')

enter image description here

Perfplot also checks that the value returned by using_sliding equals the value returned by using_alt.

Matt Timmermans' idea, "subtract position_in_a from the count for that bin" triggered this solution.

The number of elements in a bin b is the number of elements <= b.end minus the number of elements < b.start.

So you can make a starts array of the bins sorted by start, and and ends array of bins sorted by end. Then walk through all 3 arrays in step. When you advance past each x in a, advance past the starts with x < b.start and subtract position_in_a from the count for that bin. Then advance past the ends with x <= b.end and add position_in_a to the count for that bin.

Total complexity is O(N log N), dominated by sorting the start and end arrays. Walking through the 3 arrays and adjusting the counts is O(N).

In your code you are generating the array of bins already sorted, so if you can do that then you can skip the sorting step and the total complexity is O(a.length+bin_count). I would not bother even generating that array since you can easily calculate the start and end values from the index.

Something like this (?):

def sliding_count(a, nx0, nx1, window):
    bin0 = np.arange(nx0,nx1,1)
    bin1 = bin0 + window 
    count = np.zeros((nx1-nx0), dtype=int)

    for j in range(nx1-nx0):
        count[j] = np.sum(a<=bin1[j]) - np.sum(a<bin0[j])
    return count

#---- main ---------------  
nx0, nx1, window = 0, 11, 10
a = np.array([1, 5, 8, 11, 14, 19])
sliding_count(a, nx0, nx1, window)

array([3, 4, 3, 3, 4, 4, 3, 3, 3, 3, 3])

I did not check the code for nx0>0 and step>1 in bin0 = np.arange(nx0,nx1,1). So the lenght of the for-loop has to be modified for such cases.

Related