python - 对任意数量的数组的所有可能组合求和并应用限制并返回索引

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)

产生结果:

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

和索引:

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