N-D версия itertools.combinations в numpy

Я хотел бы реализовать itertools.combinations для numpy. Основываясь на этом обсуждении, у меня есть функция, которая работает для ввода 1D:

def combs(a, r):
    """
    Return successive r-length combinations of elements in the array a.
    Should produce the same output as array(list(combinations(a, r))), but 
    faster.
    """
    a = asarray(a)
    dt = dtype([('', a.dtype)]*r)
    b = fromiter(combinations(a, r), dt)
    return b.view(a.dtype).reshape(-1, r)

и вывод имеет смысл:

In [1]: list(combinations([1,2,3], 2))
Out[1]: [(1, 2), (1, 3), (2, 3)]

In [2]: array(list(combinations([1,2,3], 2)))
Out[2]: 
array([[1, 2],
       [1, 3],
       [2, 3]])

In [3]: combs([1,2,3], 2)
Out[3]: 
array([[1, 2],
       [1, 3],
       [2, 3]])

однако было бы лучше, если бы я мог расширить его до N-D входов, где дополнительные размеры просто позволяют вам быстро выполнять несколько вызовов одновременно. Итак, концептуально, если combs([1, 2, 3], 2) создает [1, 2], [1, 3], [2, 3], а combs([4, 5, 6], 2) создает [4, 5], [4, 6], [5, 6], тогда combs((1,2,3) and (4,5,6), 2) должен создавать [1, 2], [1, 3], [2, 3] and [4, 5], [4, 6], [5, 6], где "и" просто представляет собой параллельные строки или столбцы (в зависимости от того, что имеет смысл). (а также для дополнительных измерений)

Я не уверен:

  • Как заставить измерения работать логически, что согласуется с тем, как работают другие функции (например, как некоторые функции numpy имеют параметр axis= и значение по умолчанию оси 0. Так что, вероятно, ось 0 должна быть той, Комбинируя вместе, а все остальные оси представляют собой параллельные вычисления?)
  • Как заставить вышеуказанный код работать с ND (сейчас я получаю ValueError: setting an array element with a sequence.)
  • Есть ли лучший способ сделать dt = dtype([('', a.dtype)]*r)?

Ответы

Ответ 1

Вы можете использовать itertools.combinations() для создания массива индексов, а затем использовать индексирование NumPy:

import numpy as np
from itertools import combinations, chain
from scipy.misc import comb

def comb_index(n, k):
    count = comb(n, k, exact=True)
    index = np.fromiter(chain.from_iterable(combinations(range(n), k)), 
                        int, count=count*k)
    return index.reshape(-1, k)

data = np.array([[1,2,3,4,5],[10,11,12,13,14]])

idx = comb_index(5, 3)
print data[:, idx]

выход:

[[[ 1  2  3]
  [ 1  2  4]
  [ 1  2  5]
  [ 1  3  4]
  [ 1  3  5]
  [ 1  4  5]
  [ 2  3  4]
  [ 2  3  5]
  [ 2  4  5]
  [ 3  4  5]]

 [[10 11 12]
  [10 11 13]
  [10 11 14]
  [10 12 13]
  [10 12 14]
  [10 13 14]
  [11 12 13]
  [11 12 14]
  [11 13 14]
  [12 13 14]]]

Ответ 2

Не уверен, как он будет работать по производительности, но вы можете делать комбинации в индексном массиве, а затем извлекать реальные срезы массива с помощью np.take:

def combs_nd(a, r, axis=0):
    a = np.asarray(a)
    if axis < 0:
        axis += a.ndim
    indices = np.arange(a.shape[axis])
    dt = np.dtype([('', np.intp)]*r)
    indices = np.fromiter(combinations(indices, r), dt)
    indices = indices.view(np.intp).reshape(-1, r)
    return np.take(a, indices, axis=axis)

>>> combs_nd([1,2,3], 2)
array([[1, 2],
       [1, 3],
       [2, 3]])
>>> combs_nd([[1,2,3],[4,5,6]], 2, axis=1)
array([[[1, 2],
        [1, 3],
        [2, 3]],

       [[4, 5],
        [4, 6],
        [5, 6]]])