Variable number of nested for loops with fixed ran

2020-05-01 03:19发布

I'm looking for a method to have a variable number of nested for loops instead of the following code. For example if the variable n represents the number of nested for loops and n = 3, my code would be:

p = []
for i in range(26):
  for j in range(26):
    for k in range(26):
      p.append([i,j,k])

Instead, if n = 2 my code would be:

p = []
for i in range(26):
  for j in range(26):
    p.append([i,j])

I understand this can be done using recursion but I'm a bit confused as to how I'd go about doing this.

2条回答
▲ chillily
2楼-- · 2020-05-01 04:21

Something like this should work:

import itertools

n=3
fixed=26
p = list(itertools.product(range(fixed), repeat=n))

This solution uses the optimized functions of itertools, so it should be quite fast.

Mind that itertools.product returns an iterator, so one needs to transform it to get an array.

查看更多
成全新的幸福
3楼-- · 2020-05-01 04:25

It's important for one to develop the skills to reason about these problems. In this case, Python includes itertools.product but what happens the next time you need to write a behaviour specific to your program? Will there be another magical built-in function? Maybe someone else will have published a 3rd party library to solve your problem?

Below, we design product as a simple recursive function that accepts 1 or more lists.

def product (first, *rest):
  if not rest:
    for x in first:
      yield (x,)
  else:
    for p in product (*rest):
      for x in first:
        yield (x, *p)

for p in product (range(2), range(2), range(2)):
  print ('x: %d, y: %d z: %d' % p)

# x: 0, y: 0 z: 0
# x: 1, y: 0 z: 0
# x: 0, y: 1 z: 0
# x: 1, y: 1 z: 0
# x: 0, y: 0 z: 1
# x: 1, y: 0 z: 1
# x: 0, y: 1 z: 1
# x: 1, y: 1 z: 1

Assuming you want a more conventional iteration ordering, you can accomplish do so by using an auxiliary loop helper

def product (first, *rest):
  def loop (acc, first, *rest):
    if not rest:
      for x in first:
        yield (*acc, x)
    else:
      for x in first:
        yield from loop ((*acc, x), *rest)
  return loop ((), first, *rest)

for p in product (range(2), range(2), range(2)):
  print ('x: %d, y: %d z: %d' % p)

# x: 0, y: 0 z: 0
# x: 0, y: 0 z: 1
# x: 0, y: 1 z: 0
# x: 0, y: 1 z: 1
# x: 1, y: 0 z: 0
# x: 1, y: 0 z: 1
# x: 1, y: 1 z: 0
# x: 1, y: 1 z: 1
查看更多
登录 后发表回答