Tukukkk

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.

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

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.

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