Improve performance in lists

Viewed 323

I have a problem where I am trying to take a randomly ordered list and I want to know how many elements with a greater index than the current element are smaller in value than the current element.

For example:

[1,2,5,3,7,6,8,4]

should return:

[0,0,2,0,2,1,1,0]

This is the code I have that is currently working.

bribe_array = [0] * len(q)
for i in range(0, len(bribe_array)-1):
    bribe_array[i] = sum(j<q[i] for j in q[(i+1):])

This does produce the desired array but it runs slowly. What is the more pythonic way to get this accomplished?

4 Answers

We could fiddle around with the code in the question, but still it would be an O(n^2) algorithm. To truly improve the performance is not a matter of making the implementation more or less pythonic, but to use a different approach with a helper data structure.

Here's an outline for an O(n log n) solution: implement a self-balancing BST (AVL or red-black are good options), and additionally store in each node an attribute with the size of the subtree rooted in it. Now traverse the list from right to left and insert all its elements in the tree as new nodes. We also need an extra output list of the same size of the input list to keep track of the answer.

For every node we insert in the tree, we compare its key with the root. If it's greater than the value in the root, it means that it's greater than all the nodes in the left subtree, hence we need to add the size of the left subtree to the answer list at the position of the element we're trying to insert.

We keep doing this recursively and updating the size attribute in each node we visit, until we find the right place to insert the new node, and proceed to the next element in the input list. In the end the output list will contain the answer.

Another option that's much simpler than implementing a balanced BST is to adapt merge sort to count inversions and accumulate them during the process. Clearly, any single swap is an inversion so the lower-indexed element gets one count. Then during the merge traversal, simply keep track of how many elements from the right group have moved to the left and add that count for elements added to the right group.

Here's a very crude illustration :)

[1,2,5,3,7,6,8,4]

sort 1,2 | 5,3
           3,5 -> 5: 1
merge
1,2,3,5

sort 7,6 | 8,4
     6,7 -> 7: 1
           4,8 -> 8: 1
merge
4 -> 6: 1, 7: 2
4,6,7,8

merge 1,2,3,5 | 4,6,7,8
1,2,3,4 -> 1 moved
5 -> +1 -> 5: 2
6,7,8

There are several ways of speeding up your code without touching the overall computational complexity.

This is so because there are several ways of writing this very algorithm.

Let's start with your code:

def bribe_orig(q):
    bribe_array = [0] * len(q)
    for i in range(0, len(bribe_array)-1):
        bribe_array[i] = sum(j<q[i] for j in q[(i+1):])
    return bribe_array

This is of somewhat mixed style: firstly, you generate a list of zeros (which is not really needed, as you can append items on demand; secondly, the outer list uses a range() which is sub-optimal given that you would like to access a specific item multiple times and hence a local name would be faster; thirdly, you write a generator inside sum() which is also sub-optimal since it will be summing up booleans and hence perform implicit conversions all the time.

A cleaner approach would be:

def bribe(items):
    result = []
    for i, item in enumerate(items):
        partial_sum = 0
        for x in items[i + 1:]:
            if x < item:
                partial_sum += 1
        result.append(partial_sum)
    return result

This is somewhat simpler and since it does a number of things explicitly, and only performing a summation when necessary (thus skipping when you would be adding 0), it may be faster.

Another way of writing your code in a more compact way is:

def bribe_compr(items):
    return [sum(x < item for x in items[i + 1:]) for i, item in enumerate(items)]

This involves the use of generators and list comprehensions, but also the outer loop is written with enumerate() following the typical Python style.

But Python is infamously slow in raw looping, therefore when possible, vectorization can be helpful. One way of doing this (only for the inner loop) is with numpy:

import numpy as np

def bribe_np(items):
    items = np.array(items)
    return [np.sum(items[i + 1:] < item) for i, item in enumerate(items)]

Finally, it is possible to use a JIT compiler to speed up the plain Python loops using Numba:

import numba

bribe_jit = nb.jit(bribe)

As for any JIT it has some costs for the just-in-time compilation, which is eventually offset for large enough loops. Unfortunately, Numba's JIT does not support all Python code, but when it does, like in this case, it can be pretty rewarding.

Let's look at some numbers.

Consider the input generated with the following:

import numpy as np

np.random.seed(0)

n = 10
q = np.random.randint(1, n, n)

On a small-sized inputs (n = 10):

%timeit bribe_orig(q)
# 228 µs ± 3.56 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit bribe(q)
# 20.3 µs ± 814 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit bribe_compr(q)
# 216 µs ± 5.32 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit bribe_np(q)
# 133 µs ± 9.16 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit bribe_jit(q)
# 1.11 µs ± 17.9 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

On a medium-sized inputs (n = 100):

%timeit bribe_orig(q)
# 20.5 ms ± 398 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit bribe(q)
# 741 µs ± 11.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit bribe_compr(q)
# 18.9 ms ± 202 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit bribe_np(q)
# 1.22 ms ± 27.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit bribe_jit(q)
# 7.54 µs ± 165 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

On a larger inputs (n = 10000):

%timeit bribe_orig(q)
# 1.99 s ± 19.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit bribe(q)
# 60.6 ms ± 280 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit bribe_compr(q)
# 1.8 s ± 11.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit bribe_np(q)
# 12.8 ms ± 32.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit bribe_jit(q)
# 182 µs ± 2.66 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

From these results, we observe that we gain the most from substituting sum() with the explicit construct involving only Python loops. The use of comprehensions does not land you above approx. 10% improvement over your code. For larger inputs, the use of NumPy can be even faster than the explicit construct involving only Python loops. However, you will get the real deal when you use the Numba's JITed version of bribe().

You can get better performance by progressively building a sorted list going from last to first in your array. Using a binary search algorithm on the sorted list for each element in the array, you get the index at which the element will be inserted which also happens to be the number of elements that are smaller in the ones already processed.

Collecting these insertion points will give you the expected result (in reverse order).

Here's an example:

a   = [1,2,5,3,7,6,8,4]

from bisect import bisect_left
s = []
r = []
for x in reversed(a):
    p = bisect_left(s,x)
    r.append(p)
    s.insert(p,x)
r = r[::-1]

print(r) #[0,0,2,0,2,1,1]

For this example, the progression will be as follows:

step 1: x = 4,  p=0 ==>  r=[0]               s=[4]
step 2: x = 8,  p=1 ==>  r=[0,1]             s=[4,8]
step 3: x = 6,  p=1 ==>  r=[0,1,1]           s=[4,6,8]
step 4: x = 7,  p=2 ==>  r=[0,1,1,2]         s=[4,6,7,8]
step 5: x = 3,  p=0 ==>  r=[0,1,1,2,0]       s=[3,4,6,7,8]
step 6: x = 5,  p=2 ==>  r=[0,1,1,2,0,2]     s=[3,4,5,6,7,8] 
step 7: x = 2,  p=0 ==>  r=[0,1,1,2,0,2,0]   s=[2,3,4,5,6,7,8] 
step 8: x = 1,  p=0 ==>  r=[0,1,1,2,0,2,0,0] s=[1,2,3,4,5,6,7,8]

Reverse r, r = r[::-1]   r=[0,0,2,0,2,1,1,0]

You will be performing N loops (size of the array) and the binary search performs in log(i) where i is 1 to N. So, smaller than O(N*log(N)). The only caveat is the performance of s.insert(p,x) which will introduce some variability depending on the order of the original list. Overall the performance profile should be between O(N) and O(N*log(N)) with a worst case of O(n^2) when the array is already sorted.

If you only need to make your code a little faster and more concise, you could use a list comprehension (but that'll still be O(n^2) time):

r   = [sum(v<p for v in a[i+1:]) for i,p in enumerate(a)]
Related