多層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']