What is a fast, pythonic way to sum items in a list at specified indices?

Question

I have to sum up a subset of items in a list. The numbers I use to calculate the indices are given in another list. The indices may change, but the list I'm summing from does not. Both lists are of known size. This bit of code is in a nested loop in my program, and it's by far the slowest part.

I need a fast and clean way to do this in Python 3.

An obvious solution I have tried is just hardcoding the sum of all the different items. I've also tried a much cleaner solution using enumerate and summing a comprehension. The problem is that the latter is much slower than the former.

The exact index to use is 2 * i + x where i is the index in indices and x is the number in indices. (The indices list represents a set of choices between values concatenated into the lookup table.)

# sample code - the real lists are much larger

lookup = [7, 10, 1, 4, 1, 7, 9, 3, 5, 6]
indices = [0, 1, 0, 0, 1]

# hardcoded solution
s = lookup[2 * 0 + indices[0]] + lookup[2 * 1 + indices[1]] + lookup[2 * 2 + indices[2]] + lookup[2 * 3 + indices[3]] + lookup[2 * 4 + indices[4]]

# Pythonic solution with enumerate
s = sum(lookup[2 * i + x] for i, x in enumerate(indices))

I have tested both of these options using perf. The clean solution with enumerate is just over 3 times slower than the hardcoded version. The rest of the code is fairly optimized, so my whole program is almost 3 times slower if I go with the clean version.

Is there something faster I could do? Answers that require preprocessing the lookup list in some way are okay: that list is only built once but used many times.

Edit: here is a complete example of a case where hardcoding lookups to a list seems to be much faster than any alternative. The following code runs in Pypy3 in 0.27s, and the commented out slow version runs in 2.8s. (Obviously there are faster ways to do this particular task.)

from itertools import product

lookup = [1, 7, 7, 1, 2, 9, 9, 9, 2, 2, 8, 8, 9, 6, 5, 10, 3, 4, 7, 10, 1, 3, 0, 1, 7, 1, 3, 4, 2, 9]
largest_sum = 0
largest_sum_indices = []

for indices in product(list(range(0,2)), repeat=15):
    # simulate checking many different lookup lists
    for _ in range(200):
        s = lookup[2 * 0 + indices[0]] + lookup[2 * 1 + indices[1]] + lookup[2 * 2 + indices[2]] + lookup[2 * 3 + indices[3]] + lookup[2 * 4 + indices[4]] + lookup[2 * 5 + indices[5]] + lookup[2 * 6 + indices[6]] + lookup[2 * 7 + indices[7]] + lookup[2 * 8 + indices[8]] + lookup[2 * 9 + indices[9]] + lookup[2 * 10 + indices[10]] + lookup[2 * 11 + indices[11]] + lookup[2 * 12 + indices[12]] + lookup[2 * 13 + indices[13]] + lookup[2 * 14 + indices[14]]
        # clean method is too slow
        #s = sum(lookup[i * 2 + x] for i,x in enumerate(indices))
        if s > largest_sum:
            largest_sum = s
            largest_sum_indices = indices

print(largest_sum)
print(largest_sum_indices)

Answer 1

I think that you really need to re-evaluate your algorithm instead of attempting to squeeze out cycles from innocent code if the edit is indicative of what you are really trying to do.

Let's start by describing in words what your code does. You have a lookup array of size 2N. You are assigning a bit (0 or 1) to indicate which element you will select from each successive pair, and add up the N selected elements. By going through every bit combination from 0 to 2**N-1, you hope to find the maximum sum of N elements.

I would posit that simply inspecting the N pairs of elements in a single pass will give you the correct sum and indices, if you still want those. You can do this in N steps, not N * 2**N.

Here is a really basic, totally unoptimized, solution that I bet will scale better than the one in the question:

lookup = ...
N = len(lookup) // 2
largest_sum = 0
largest_sum_indices = [0] * N
for i in range(N):
    if lookup[2 * i + 1] > lookup[2 * i]:
        largest_sum_indices[i] = 1
        largest_sum += lookup[2 * i + 1]
    else:
        largest_sum += lookup[2 * i]

Notice that I don't call any functions besides len and just use a basic for loop.

Here is an optimized numpy version:

import numpy as np

lookup = np.array(...)
largest_sum_indices = np.argmax(lookup.reshape(lookup.size // 2, 2), axis=1)
largest_sum = lookup[2 * np.arange(largest_sum_indices.size) + largest_sum_indices]

While your own tests will show you which algorithm works best for you, keep in mind that either of the options here could handily process a few million elements without making the user too antsy, while somerhing that scales as O(N * 2**N) would take longer than the heat death of many universes. Any time you use product, there is a good chance of a better solution being available.

Answer 2

You can use the function itemgetter() for fast lookup:

from operator import itemgetter
from itertools import count

lookup = [7, 10, 1, 4, 1, 7, 9, 3, 5, 6]
indices = [0, 1, 0, 0, 1]

sum(itemgetter(*[i + j for i, j in zip(count(step=2), indices)])(lookup))
# 27

Answer 3

This seems like a good task for numpy. It allows you to vectorize many of the operations you are doing, running C implementations under the hood.

import numpy as np

lookup = np.array([7, 10, 1, 4, 1, 7, 9, 3, 5, 6], dtype=np.int)
indices = np.array([0, 1, 0, 0, 1], dtype=np.int)

real_indices = np.arange(0, 2 * indices.size, 2) + indices
s = lookup[real_indices].sum()