Find all subarrays of fixed length with a given ranking

Viewed 1074

I have an array of numbers, for example:

A = [1, 5, 2, 4, 3]

and an array that determines a rank, for example:

B = [0, 2, 1]

My goal is to find all the subarrays of A that "obey" the rank B. If a subarray obeys the rank, that means that the i-th smallest element of the subarray must have B[i] as its (subarray) index. So for a subarray to match, the smallest element within it must be in position 0, the second smallest must be in position 2, and the biggest element must be in position 1.

So for example here, there are two subarrays of A that match the ranking: [1, 5, 2] (because A[0] < A[2] < A[1]) and [2, 4, 3].

So far, I've managed to find a solution that has an O(mn) (m is len(A) and n is len(B)) time complexity, it iterates over all the subarrays of length 3 and verifies if they are correctly ordered:

A = [1, 5, 2, 4, 3]
B = [0, 2, 1]
m = len(A)
n = len(B)
for i in range(m - n + 1):
    current_subarray = A[i:i + n]
    # we now do n - 1 comparaisons to check whether the subarray is correctly ordered.
    for B_index in range(n - 1):
        if current_subarray[B[B_index]] > current_subarray[B[B_index + 1]]:
            break
    else:
        print("Subarray found:", current_subarray)

Result:

Subarray found: [1, 5, 2]
Subarray found: [2, 4, 3]

This works, but I was wondering if there was a better optimized algorithm (better than O(mn)) to fulfill this task.

4 Answers

Instead of iterating over B to compare ranks, you could use scipy.stats.rankdata to get the ranks directly:

from scipy.stats import rankdata

A = [1, 5, 2, 4, 3]
B = [0, 2, 1]

m = len(A)
n = len(B)

for i in range(m - n + 1):
    current_subarray = A[i:i + n]

    ranked_numbers = (rankdata(current_subarray).astype(int) - 1).tolist()
    if ranked_numbers == B:
        print("Subarray found:", current_subarray)

# Subarray found: [1, 5, 2]
# Subarray found: [2, 4, 3]

Note: rankdata() starts ranks from 1 instead of 0, which is why the above minuses 1 from every element in the numpy array.

Here's a numpy solution based on some linear algebra.

First convert B to a basis:

import numpy as np
A = [1, 5, 2, 4, 3]
B = [0, 2, 1]

b = np.eye(len(B))[B]
print(b)
#array([[1, 0, 0],
#       [0, 0, 1],
#       [0, 1, 0]])

Now we can go through each subarray of A and project it into this space. If the resulting vector is sorted, that means the subarray followed the ranking.

for i in range(0, (len(A) - len(B))+1):
    a = np.array(A[i:i+len(B)])
    if (np.diff(a.dot(b))>0).all():
        print(a)
#[1 5 2]
#[2 4 3]

I'm not a numpy expert, so there may be a way to optimize this further and eliminate the loop.


Update, here's a cleaner version:

def get_ranked_subarrays(A, B):
    m = len(A)
    n = len(B)
    b = np.eye(n)[B]
    a = np.array([A[i:i+n] for i in range(0, m - n+1)])
    return a[(np.diff(a.dot(b))>0).all(1)].tolist()

A = [1, 5, 2, 4, 3]
B = [0, 2, 1]
get_ranked_subarrays(A, B)
#[[1, 5, 2], [2, 4, 3]]

Performance Results:

Your solution is very good for small n, but the numpy solution outperforms as the size of A grows large:

Here's your code which I turned into a function that returns the desired subarrays (instead of printing):

def get_ranked_subarrays_op(A, B):
    m = len(A)
    n = len(B)
    out = []
    for i in range(m - n + 1):
        current_subarray = A[i:i + n]
        # we now do n - 1 comparisons to check whether the subarray is correctly ordered.
        for B_index in range(n - 1):
            if current_subarray[B[B_index]] > current_subarray[B[B_index + 1]]:
                break
        else:
            out.append(current_subarray)
    return out

Timing results for a large random A:

array_size = 1000000
A = np.random.randint(low=0, high=10, size=array_size)
B = [0, 2, 1]

%%timeit
get_ranked_subarrays_op(A, B)
#1 loop, best of 3: 1.57 s per loop

%%timeit
get_ranked_subarrays(A, B)
#1 loop, best of 3: 890 ms per loop

However, if m also grows large, your solution is slightly better because of the short circuiting (the likelihood of a short circuit grows large for large m). Here is the timing results of we let m be 100.

array_size = 1000000
basis_size = 100
A = np.random.randint(low=0, high=10, size=array_size)
B = range(basis_size)
np.random.shuffle(B)

%%timeit
get_ranked_subarrays_op(A, B)
#1 loop, best of 3: 1.9 s per loop

%%timeit
get_ranked_subarrays(A, B)
#1 loop, best of 3: 2.79 s per loop

You can loop over A and check the resulting subarrays:

A, B = [1, 5, 2, 4, 3], [0, 2, 1]
def results(a, b):
   _l = len(b)
   for c in range(len(a)-_l+1):
     _r = a[c:c+_l]
     new_r = [_r[i] for i in b]
     if all(new_r[i] < new_r[i+1] for i in range(len(new_r)-1)):
       yield _r

print(list(results(A, B)))

Output:

[[1, 5, 2], [2, 4, 3]]

At the very least, we could could rule out candidate windows much more quickly by considering the (binary) relation of neighbouring elements, which could allow for parallel examination. Call less than 0 and greater than 1. Then:

A = [1,  5,  2,  4,  3]
A'=   [0,  1,  0,  1]

B'=   [0,  1]
B = [0,  2,  1]

Clearly any candidate must match the relation sequence. Also note that the only type of section of B that could admit overlap is an ascending or descending sequence (means we can possibly skip ahead a priori if a match is found).

Related