多层for循环转递归/迭代

使用场景

枚举组合:

问题是这样的.

有 n 个列表,分别从每个列表中取出一个元素,一共有多少种组合?

例如:

a = ['a1','a2']
b = ['b1','b2','b3']

组合结果为:

[
  ('a1','b1'),
  ('a1','b2'),
  ('a1','b3'),
  ('a2','b1'),
  ('a2','b2'),
  ('a2','b3')
]

待组合的列表只有两个

这种情况就是简单的遍历:

a = ['a1','a2']
b = ['b1','b2','b3']

res = []
for i in a:
  for j in b:
    res.append((i,j)])

print(res)

扩展为 n 个

如果还用for循环嵌套,代码就是这样的

a = ['a1','a2']
b = ['b1','b2','b3']

res = []
for i in a:
  for j in b:
    for k in c:
        ...
           ...

如果是n层的话,这样的代码是无法表达的.

我们可以先将第一个和第二个组合,再拿组合出的结果和第三个组合,依次类推...

如下如所示:

用代码表示如下:

迭代

def merge(i,j):
  """
  i = "a"
  j = ("b","c")
  return: ("a","b","c")
  """
  res = []
  for p in (i,j):
    if isinstance(p,tuple):
      res.extend(p)
    else:
      res.append(p)
  return tuple(res)

def combineN(*args):
  target = args[0]
  for li in args[1:]:
    tmp = []
    for i in target:
      for j in li:
        tmp.append(merge(i,j))
    target = tmp
  return target

递归

def merge(i,j):
  """
  i = "a"
  j = ("b","c")
  return: ("a","b","c")
  """
  res = []
  for p in (i,j):
    if isinstance(p,tuple):
      res.extend(p)
    else:
      res.append(p)
  return tuple(res)

def combine2(a, b):
    res = []
    for i in a:
        for j in b:
            res.append(merge(i,j))
    return res

def combineNRecursion(*args):
  if len(args) == 2:
    return combine2(*args)

  return combine2(args[0],combineNRecursion(*args[1:]))

通用的多层 for 循环转迭代

上面用到的迭代方法是针对具体问题分析得来的,那么有没有一种通用的转换方案呢? 答案是肯定的.

def combineN(*li):
    res = []
    # 相当于最内层循环执行的次数.
    total_times = reduce(lambda x, y: x*y, [len(item) for item in li])
    n = 0
    while n < total_times:
        tmp = n
        tem_res = []
        for i in range(len(li)):
            # 余数就是参与计算的元素的下标,商用于计算下一个列表参与元素的下标.
            tmp, cur = divmod(tmp, len(li[i]))
            tem_res.append(li[i][cur])
        res.append(tem_res)
        n += 1
    return res

res = combineN(["a1","a2"], ["b1", "b2"], ["c1", "c2"])
for i in res:
    print(i)

输出结果如下:

['a1', 'b1', 'c1']
['a2', 'b1', 'c1']
['a1', 'b2', 'c1']
['a2', 'b2', 'c1']
['a1', 'b1', 'c2']
['a2', 'b1', 'c2']
['a1', 'b2', 'c2']
['a2', 'b2', 'c2']