Суммируя все возможные комбинации произвольного числа массивов и применяя ограничения и возвращаемые индексы

Question

Это модификация this question, в которой я хотел бы вернуть индексы элементов массивов в дополнение к самим элементам. Я успешно изменил arraysums(), arraysums_recursive(), но я боюсь arraysums_recursive_anyvals(). Вот подробности:

Я изменил arraysums():

def arraysums(arrays,lower,upper): 
    products = itertools.product(*arrays) 
    result = list() 

    indices = itertools.product(*[np.arange(len(arr)) for arr in arrays]) 
    index = list() 

    for n,k in zip(products,indices): 
     s = sum(n) 
     if lower <= s <= upper: 
      result.append(n) 
      index.append(k)     
    return result,index 

Теперь возвращает элементы и индексы элементов:

N = 8 
a = np.arange(N) 
b = np.arange(N)-N/2  
arraysums((a,b),lower=5,upper=6) 


([(2, 3), 
    (3, 2), 
    (3, 3), 
    (4, 1), 
    (4, 2), 
    (5, 0), 
    (5, 1), 
    (6, -1), 
    (6, 0), 
    (7, -2), 
    (7, -1)], 
[(2, 7), 
    (3, 6), 
    (3, 7), 
    (4, 5), 
    (4, 6), 
    (5, 4), 
    (5, 5), 
    (6, 3), 
    (6, 4), 
    (7, 2), 
    (7, 3)]) 

Я также изменил @ рекурсивного решения unutbu, который также возвращает такой же результат, как arraysums():

def arraysums_recursive(arrays, lower, upper): 
    if len(arrays) <= 1: 
     result = [(item,) for item in arrays[0] if lower <= item <= upper] 
     index = [] # this needs to be fixed 
    else: 
     result = [] 
     index = [] 
     for item in arrays[0]: 
      subarrays = [[item2 for item2 in arr if item2 <= upper-item] 
         for arr in arrays[1:]] 
      result.extend(
       [(item,)+tup for tup in arraysums(
        subarrays, lower-item, upper-item)[0]]) 
      index.extend(
       [(item,)+tup for tup in arraysums(
        subarrays, lower-item, upper-item)[1]]) 

    return result,index 

Наконец, я изменил arraysums_recursive_anyvals(), но я не могу понять, почему он не возвращает индексы:

def arraysums_recursive_anyvals(arrays, lower, upper): 
    if len(arrays) <= 1: 
     result = [(item,) for item in arrays[0] if lower <= item <= upper] 
     index = [] # this needs to be fixed 
    else: 
     minval = min(item for arr in arrays for item in arr) 
     # Subtract minval from arrays to guarantee all the values are positive 
     arrays = [[item-minval for item in arr] for arr in arrays] 
     # Adjust the lower and upper bounds accordingly 
     lower -= minval*len(arrays) 
     upper -= minval*len(arrays) 

     result = [] 
     index = [] 
     for item in arrays[0]: 
      subarrays = [[item2 for item2 in arr if item2 <= upper-item] 
         for arr in arrays[1:]] 
      if min(len(arr) for arr in subarrays) == 0: 
       continue 
      result.extend(
       [(item,)+tup for tup in arraysums_recursive(
        subarrays, lower-item, upper-item)[0]]) 
      index.extend(
       [(item,)+tup for tup in arraysums_recursive(
        subarrays, lower-item, upper-item)[1]]) 

     # Readjust the result by adding back minval 
     result = [tuple([item+minval for item in tup]) for tup in result] 
    return result,index 

Результаты:

arraysums_recursive_anyvals((a,b),lower=5,upper=6) 

([(2, 3), 
    (3, 2), 
    (3, 3), 
    (4, 1), 
    (4, 2), 
    (5, 0), 
    (5, 1), 
    (6, -1), 
    (6, 0), 
    (7, -2), 
    (7, -1)], 
[]) 

Answer 1

Ключевой особенностью arraysums_recursive является то, что он выбрасывает значения который не может внести свой вклад в результат:

subarrays = [[item2 for item2 in arr if item2 <= upper-item] 
       for arr in arrays[1:]] 

Хотя бросать вещи усложняет запись индексов, это не так уж трудно. Во-первых, в arraysums_recursive расширить arrays включить индекс, а также значение элемента:

def arraysums_recursive(arrays, lower, upper): 
    arrays = [[(i, item) for i, item in enumerate(arr)] for arr in arrays] 
    ... 
    index, result = zip(*arraysums_recursive_all_positive(arrays, lower, upper)) 
    return result, index 

Теперь перепишем arraysums_recursive_all_positive обрабатывать arrays, которые состоят из списка списков (index, item)кортежей.

---

def arraysums_recursive(arrays, lower, upper): 
    arrays = [[(i, item) for i, item in enumerate(arr)] for arr in arrays] 
    minval = min(item for arr in arrays for i, item in arr) 
    # Subtract minval from arrays to guarantee all the values are positive 
    arrays = [[(i, item-minval) for i, item in arr] for arr in arrays] 
    # Adjust the lower and upper bounds accordingly 
    lower -= minval*len(arrays) 
    upper -= minval*len(arrays) 
    index, result = zip(*arraysums_recursive_all_positive(arrays, lower, upper)) 
    # Readjust the result by adding back minval 
    result = [tuple([item+minval for item in tup]) for tup in result] 
    return result, index 

def arraysums_recursive_all_positive(arrays, lower, upper): 
    # Assumes all values in arrays are positive 
    if len(arrays) <= 1: 
     result = [((i,), (item,)) for i, item in arrays[0] if lower <= item <= upper] 
    else: 
     result = [] 
     for i, item in arrays[0]: 
      subarrays = [[(i, item2) for i, item2 in arr if item2 <= upper-item] 
         for arr in arrays[1:]] 
      if min(len(arr) for arr in subarrays) == 0: 
       continue 
      result.extend(
       [((i,)+i_tup, (item,)+item_tup) for i_tup, item_tup in 
       arraysums_recursive_all_positive(subarrays, lower-item, upper-item)]) 
    return result 

---

import numpy as np 
N = 8 
a = np.arange(N) 
b = np.arange(N)-N/2  
result, index = arraysums_recursive((a,b),lower=5,upper=6) 

дает result:

[(2.0, 3.0), 
(3.0, 2.0), 
(3.0, 3.0), 
(4.0, 1.0), 
(4.0, 2.0), 
(5.0, 0.0), 
(5.0, 1.0), 
(6.0, -1.0), 
(6.0, 0.0), 
(7.0, -2.0), 
(7.0, -1.0)] 

и index:

((2, 7), 
(3, 6), 
(3, 7), 
(4, 5), 
(4, 6), 
(5, 4), 
(5, 5), 
(6, 3), 
(6, 4), 
(7, 2), 
(7, 3))