Generar una matriz numpy con todas las combinaciones de números que sumn menos que un número dado

Hay varios ejemplos elegantes de usar numpy en Python para generar arreglos de todas las combinaciones. Por ejemplo, la respuesta aquí: Usar numpy para construir una matriz de todas las combinaciones de dos matrices .

Ahora suponga que hay una restricción adicional, es decir, la sum de todos los números no puede sumr más que una constante K dada. Usando un generador e itertools.product , para un ejemplo con K=3 donde queremos las combinaciones de tres variables con rangos 0-1,0-3, y 0-2 podemos hacerlo de la siguiente manera:

 from itertools import product K = 3 maxRange = np.array([1,3,2]) states = np.array([i for i in product(*(range(i+1) for i in maxRange)) if sum(i)<=K]) 

que devuelve

 array([[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 1, 0], [0, 1, 1], [0, 1, 2], [0, 2, 0], [0, 2, 1], [0, 3, 0], [1, 0, 0], [1, 0, 1], [1, 0, 2], [1, 1, 0], [1, 1, 1], [1, 2, 0]]) 

En principio, el enfoque de https://stackoverflow.com/a/25655090/1479342 se puede usar para generar todas las combinaciones posibles sin la restricción y luego seleccionar el subconjunto de combinaciones que sumn menos de K Sin embargo, ese enfoque genera muchas más combinaciones de las necesarias, especialmente si K es relativamente pequeño en comparación con la sum(maxRange) .

Debe haber una manera de hacerlo más rápido y con un menor uso de memoria. ¿Cómo se puede lograr esto utilizando un enfoque vectorizado (por ejemplo, usando np.indices )?

Editado

  1. Para completar, aquí estoy agregando el código del OP:

     def partition0(max_range, S): K = len(max_range) return np.array([i for i in itertools.product(*(range(i+1) for i in max_range)) if sum(i)<=S]) 
  2. El primer enfoque es puro np.indices . Es rápido para una entrada pequeña pero consume mucha memoria (OP ya señaló que no es lo que quería decir).

     def partition1(max_range, S): max_range = np.asarray(max_range, dtype = int) a = np.indices(max_range + 1) b = a.sum(axis = 0) <= S return (a[:,b].T) 
  3. El enfoque recurrente parece ser mucho mejor que los anteriores:

     def partition2(max_range, max_sum): max_range = np.asarray(max_range, dtype = int).ravel() if(max_range.size == 1): return np.arange(min(max_range[0],max_sum) + 1, dtype = int).reshape(-1,1) P = partition2(max_range[1:], max_sum) # S[i] is the largest summand we can place in front of P[i] S = np.minimum(max_sum - P.sum(axis = 1), max_range[0]) offset, sz = 0, S.size out = np.empty(shape = (sz + S.sum(), P.shape[1]+1), dtype = int) out[:sz,0] = 0 out[:sz,1:] = P for i in range(1, max_range[0]+1): ind, = np.nonzero(S) offset, sz = offset + sz, ind.size out[offset:offset+sz, 0] = i out[offset:offset+sz, 1:] = P[ind] S[ind] -= 1 return out 
  4. Después de un breve pensamiento, fui capaz de llevarlo un poco más lejos. Si conocemos de antemano el número de particiones posibles, podemos asignar suficiente memoria a la vez. (Es algo similar a cartesian en un hilo ya enlazado ).

    Primero, necesitamos una función que cuente particiones.

     def number_of_partitions(max_range, max_sum): ''' Returns an array arr of the same shape as max_range, where arr[j] = number of admissible partitions for j summands bounded by max_range[j:] and with sum <= max_sum ''' M = max_sum + 1 N = len(max_range) arr = np.zeros(shape=(M,N), dtype = int) arr[:,-1] = np.where(np.arange(M) <= min(max_range[-1], max_sum), 1, 0) for i in range(N-2,-1,-1): for j in range(max_range[i]+1): arr[j:,i] += arr[:Mj,i+1] return arr.sum(axis = 0) 

    La función principal:

     def partition3(max_range, max_sum, out = None, n_part = None): if out is None: max_range = np.asarray(max_range, dtype = int).ravel() n_part = number_of_partitions(max_range, max_sum) out = np.zeros(shape = (n_part[0], max_range.size), dtype = int) if(max_range.size == 1): out[:] = np.arange(min(max_range[0],max_sum) + 1, dtype = int).reshape(-1,1) return out P = partition3(max_range[1:], max_sum, out=out[:n_part[1],1:], n_part = n_part[1:]) # P is now a useful reference S = np.minimum(max_sum - P.sum(axis = 1), max_range[0]) offset, sz = 0, S.size out[:sz,0] = 0 for i in range(1, max_range[0]+1): ind, = np.nonzero(S) offset, sz = offset + sz, ind.size out[offset:offset+sz, 0] = i out[offset:offset+sz, 1:] = P[ind] S[ind] -= 1 return out 
  5. Algunas pruebas:

     max_range = [3, 4, 6, 3, 4, 6, 3, 4, 6] for f in [partition0, partition1, partition2, partition3]: print(f.__name__ + ':') for max_sum in [5, 15, 25]: print('Sum %2d: ' % max_sum, end = '') %timeit f(max_range, max_sum) print() partition0: Sum 5: 1 loops, best of 3: 859 ms per loop Sum 15: 1 loops, best of 3: 1.39 s per loop Sum 25: 1 loops, best of 3: 3.18 s per loop partition1: Sum 5: 10 loops, best of 3: 176 ms per loop Sum 15: 1 loops, best of 3: 224 ms per loop Sum 25: 1 loops, best of 3: 403 ms per loop partition2: Sum 5: 1000 loops, best of 3: 809 µs per loop Sum 15: 10 loops, best of 3: 62.5 ms per loop Sum 25: 1 loops, best of 3: 262 ms per loop partition3: Sum 5: 1000 loops, best of 3: 853 µs per loop Sum 15: 10 loops, best of 3: 59.1 ms per loop Sum 25: 1 loops, best of 3: 249 ms per loop 

    Y algo más grande:

     %timeit partition0([3,6] * 5, 20) 1 loops, best of 3: 11.9 s per loop %timeit partition1([3,6] * 5, 20) The slowest run took 12.68 times longer than the fastest. This could mean that an intermediate result is being cached 1 loops, best of 3: 2.33 s per loop # MemoryError in another test %timeit partition2([3,6] * 5, 20) 1 loops, best of 3: 877 ms per loop %timeit partition3([3,6] * 5, 20) 1 loops, best of 3: 739 ms per loop 

No sé cuál es el enfoque de numpy , pero aquí hay una solución razonablemente limpia. Sea A una matriz de enteros y sea k un número que se le da como entrada.

Comience con una matriz vacía B ; mantenga la sum de la matriz B en una variable s (inicialmente establecida en cero). Aplicar el siguiente procedimiento:

  • si la sum s de la matriz B es menor que k , entonces (i) agréguela a la colección, (ii) y para cada elemento de la matriz original A , agregue ese elemento a B y actualice s , (iii) elimínela de A y (iv) aplicar recursivamente el procedimiento; (iv) cuando vuelva la llamada, agregue el elemento a A y actualice s ; de lo contrario no hacer nada.

Este enfoque de abajo hacia arriba elimina las twigs no válidas desde el principio y solo visita los subconjuntos necesarios (es decir, casi solo los subconjuntos que sumn menos de k ).