Сопоставить элемент в многомерном массиве с его индексом

Question

Я использую функцию get_tuples(length, total) из здесь , чтобы сгенерировать массив всех кортежей заданной длины и суммы, пример и функция показаны ниже. После того, как я создал массив, мне нужно найти способ вернуть индексы заданного числа элементов в массиве. Я смог сделать это, используя .index(), изменив массив в список, как показано ниже. Однако это решение или другое решение, которое также основано на поиске (например, с использованием np.where), занимает много времени для поиска индексов. Поскольку все элементы в массиве (массив s в примере) различны, мне было интересно, можем ли мы построить взаимно-однозначное отображение, то есть функцию, такую, что для данного элемента в массиве он возвращает индексэлемент путем некоторого сложения и умножения на значения этого элемента. Есть идеи, если это возможно? Спасибо!

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

Вот формула, которая вычисляет индекс на основе одного кортежа, то есть он не должен видеть полный массив. Для вычисления индекса N-кортежа необходимо оценить N-1 биномиальные коэффициенты. Следующая реализация (частично) векторизована, она принимает ND-массивы, но кортежи должны быть в последнем измерении.

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

Пример выполнения:

all ok
[  0   7 100   5  45]

Как получить формулу:

1. Используйте звездочки и столбцы , чтобы выразить полное число разделов #part(N,k) (N - общее, k - длина) в виде одного биномиального коэффициента (N + k - 1) choose (k - 1).

2. Подсчет задом наперед: нетрудно проверить, что после i-й полной итерации внешнего цикла генератора ОП точно не было перечислено #part(N-i,k). В самом деле, все, что осталось, это разбиения p1 + p2 + ... = N с p1> = i;мы можем написать p1 = q1 + i так, что q1 + p2 + ... = Ni, и этот последний раздел не имеет ограничений, поэтому мы можем использовать 1. для подсчета.

Answer 2

Вы можете использовать бинарный поиск, чтобы сделать поиск намного быстрее.

Бинарный поиск делает поиск O (log (n)), а не O (n) (используя индекс)

Нам не нужно сортировать кортежи, поскольку они уже отсортированы генератором

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

Выход

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)

Производительность

Сценарий 1 - Кортежи уже вычислены, и мы просто хотим найти индекс определенных кортежей

Например, 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]]

Двоичный поиск

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

Вычисляет индекс на основе одного кортежа (любезность @PaulPanzer подход)

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

В этом сценариидвоичный поиск выполняется в ~ 8 раз быстрее, если кортежи уже были предварительно вычислены.

Сценарий 2 - кортежи не были предварительно вычислены.

%%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 в этом сценарии совпадает по времени (92,97 мкс)

=> Подход @PaulPanzer ~ в 29 раз быстрее, когда кортежи ненеобходимо вычислить

Сценарий 3 - Большое количество индексов (@PJORR) Сгенерировано большое количество случайных индексов

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

В этом случаемы @PaulPanzer на 53% быстрее