Map an element in a multi-dimension array to its index

Question

I am using the function get_tuples(length, total) from here to generate an array of all tuples of given length and sum, an example and the function are shown below. After I have created the array I need to find a way to return the indices of a given number of elements in the array. I was able to do that using .index() by changing the array to a list, as shown below. However, this solution or another solution that is also based on searching (for example using np.where) takes a lot of time to find the indices. Since all elements in the array (array s in the example) are different, I was wondering if we can construct a one-to-one mapping, i.e., a function such that given the element in the array it returns the index of the element by doing some addition and multiplication on the values of this element. Any ideas if that is possible? Thanks!

import numpy as np

def get_tuples(length, total):
    if length == 1:
        yield (total,)
        return

    for i in range(total + 1):
        for t in get_tuples(length - 1, total - i):
            yield (i,) + t
#example
s = np.array(list(get_tuples(4, 20)))

# array s
In [1]: s
Out[1]: 
array([[ 0,  0,  0, 20],
       [ 0,  0,  1, 19],
       [ 0,  0,  2, 18],
       ...,
       [19,  0,  1,  0],
       [19,  1,  0,  0],
       [20,  0,  0,  0]])

#example of element to find the index for. (Note in reality this is 1000+ elements)
elements_to_find =np.array([[ 0,  0,  0, 20],
                            [ 0,  0,  7, 13],
                            [ 0,  5,  5, 10],
                            [ 0,  0,  5, 15],
                            [ 0,  2,  4, 14]])
#change array to list
s_list = s.tolist()

#find the indices
indx=[s_list.index(i) for i in elements_to_find.tolist()]

#output
In [2]: indx
Out[2]: [0, 7, 100, 5, 45]

Answer 1

You can use binary search to make the search a lot faster.

Binary search makes the search O(log(n)) rather than O(n) (using Index)

We do not need to sort the tuples since they are already sorted by the generator

import bisect

def get_tuples(length, total):
  " Generates tuples "
  if length == 1:
    yield (total,)
    return

  yield from ((i,) + t for i in range(total + 1) for t in get_tuples(length - 1, total - i))

def find_indexes(x, indexes):
   if len(indexes) > 100:
        # Faster to generate all indexes when we have a large
        # number to check
        d = dict(zip(x, range(len(x))))
        return [d[tuple(i)] for i in indexes]
    else:
        return [bisect.bisect_left(x, tuple(i)) for i in indexes]

# Generate tuples (in this case 4, 20)
x = list(get_tuples(4, 20))

# Tuples are generated in sorted order [(0,0,0,20), ...(20,0,0,0)]
# which allows binary search to be used
indexes = [[ 0,  0,  0, 20],
           [ 0,  0,  7, 13],
           [ 0,  5,  5, 10],
           [ 0,  0,  5, 15],
           [ 0,  2,  4, 14]]

y = find_indexes(x, indexes)
print('Found indexes:', *y)
print('Indexes & Tuples:')
for i in y:
  print(i, x[i])

Output

Found indexes: 0 7 100 5 45
Indexes & Tuples:
0 (0, 0, 0, 20)
7 (0, 0, 7, 13)
100 (0, 5, 5, 10)
5 (0, 0, 5, 15)
45 (0, 2, 4, 14)

Performance

Scenario 1--Tuples already computed and we just want to find the index of certain tuples

For instance x = list(get_tuples(4, 20)) has already been perform.

Search for

indexes = [[ 0,  0,  0, 20],
           [ 0,  0,  7, 13],
           [ 0,  5,  5, 10],
           [ 0,  0,  5, 15],
           [ 0,  2,  4, 14]]

Binary Search

%timeit find_indexes(x, indexes)
100000 loops, best of 3: 11.2 µs per loop

Calculates the index based on the tuple alone (courtesy @PaulPanzer approach)

%timeit get_idx(indexes)
10000 loops, best of 3: 92.7 µs per loop

In this scenario, binary search is ~8x faster when tuples have already been pre-computed.

Scenario 2--the tuples have not been pre-computed.

%%timeit
import bisect

def find_indexes(x, t):
    " finds the index of each tuple in list t (assumes x is sorted) "
    return [bisect.bisect_left(x, tuple(i)) for i in t]

# Generate tuples (in this case 4, 20)
x = list(get_tuples(4, 20))

indexes = [[ 0,  0,  0, 20],
           [ 0,  0,  7, 13],
           [ 0,  5,  5, 10],
           [ 0,  0,  5, 15],
           [ 0,  2,  4, 14]]

y = find_indexes(x, indexes)

100 loops, best of 3: 2.69 ms per loop

@PaulPanzer approach is the same timing in this scenario (92.97 us)

=> @PaulPanzer approach ~29 times faster when the tuples don't have to be computed

Scenario 3--Large number of indexes (@PJORR) A large number of random indexes is generated

x = list(get_tuples(4, 20))
xnp = np.array(x)
indices = xnp[np.random.randint(0,len(xnp), 2000)]
indexes = indices.tolist()
%timeit find_indexes(x, indexes)
#Result: 1000 loops, best of 3: 1.1 ms per loop
%timeit get_idx(indices)
#Result: 1000 loops, best of 3: 716 µs per loop

In this case, we are @PaulPanzer is 53% faster

Answer 2

Here is a formula that calculates the index based on the tuple alone, i.e. it needn't see the full array. To compute the index of an N-tuple it needs to evaluate N-1 binomial coefficients. The following implementation is (part-) vectorized, it accepts ND-arrays but the tuples must be in the last dimension.

import numpy as np
from scipy.special import comb

# unfortunately, comb with option exact=True is not vectorized
def bc(N,k):
    return np.round(comb(N,k)).astype(int)

def get_idx(s):
    N = s.shape[-1] - 1
    R = np.arange(1,N)
    ps = s[...,::-1].cumsum(-1)
    B = bc(ps[...,1:-1]+R,1+R)
    return bc(ps[...,-1]+N,N) - ps[...,0] - 1 - B.sum(-1)

# OP's generator
def get_tuples(length, total):
    if length == 1:
        yield (total,)
        return

    for i in range(total + 1):
        for t in get_tuples(length - 1, total - i):
            yield (i,) + t
#example
s = np.array(list(get_tuples(4, 20)))

# compute each index
r = get_idx(s)

# expected: 0,1,2,3,...
assert (r == np.arange(len(r))).all()
print("all ok")

#example of element to find the index for. (Note in reality this is 1000+ elements)
elements_to_find =np.array([[ 0,  0,  0, 20],
                            [ 0,  0,  7, 13],
                            [ 0,  5,  5, 10],
                            [ 0,  0,  5, 15],
                            [ 0,  2,  4, 14]])

print(get_idx(elements_to_find))

Sample run:

all ok
[  0   7 100   5  45]

How to derive formula:

1. Use stars and bars to express the full partition count #part(N,k) (N is total, k is length) as a single binomial coefficient (N + k - 1) choose (k - 1).

2. Count back-to-front: It is not hard to verify that after the i-th full iteration of the outer loop of OP's generator exactly #part(N-i,k) have not yet been enumerated. Indeed, what's left are all partitions p1+p2+... = N with p1>=i; we can write p1=q1+i such that q1+p2+... = N-i and this latter partition is constraint-free so we can use 1. to count.