0

I'm trying to write a function that outputs a dictionary like this:

0: [0,0,0]
1: [0,0,1]
2: [0,0,2]
...
n: [0,0,n]
n+1: [0,1,0]
n+2: [0,1,1]
...
n^3: [n,n,n]

This example is when the list is 3 elements long. I'd like to generalize this to lists of arbitrary length.

Here's some code I have tried already which works for 3 element long lists. Clearly can extend this to arbitrary lengths by adding more for loops, but I'm sure there is a better way.

tricount = 0
wgtlist = [0.0, 10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0]
for i in range(0, len(wgtlist)):
    for j in range (0, len(wgtlist)):
        for k in range (0, len(wgtlist)):
            trilist = []
            trilist.append(wgtlist[i])
            trilist.append(wgtlist[j])
            trilist.append(wgtlist[k])
            trials[tricount] = trilist
            tricount = tricount+1
print (trials)

This outputs:

{0: [0.0, 0.0, 0.0], 1: [0.0, 0.0, 10.0], 2: [0.0, 0.0, 20.0], 3: [0.0, 0.0, 30.0], 4: [0.0, 0.0, 40.0], 5: [0.0, 0.0, 50.0], 6: [0.0, 0.0, 60.0], 7: [0.0, 0.0, 70.0], 8: [0.0, 0.0, 80.0], 9: [0.0, 0.0, 90.0], 10: [0.0, 0.0, 100.0], 11: [0.0, 10.0, 0.0], 12: [0.0, 10.0, 10.0], 13: [0.0, 10.0, 20.0].....

1323: [100.0, 100.0, 30.0], 1324: [100.0, 100.0, 40.0], 1325: [100.0, 100.0, 50.0], 1326: [100.0, 100.0, 60.0], 1327: [100.0, 100.0, 70.0], 1328: [100.0, 100.0, 80.0], 1329: [100.0, 100.0, 90.0], 1330: [100.0, 100.0, 100.0]}

(obviously truncated output)

1 Answers1

1

You can use itertools.product on the list with a repeat of the number of items you would like in the output tuples. Since the keys of your expected output dict are simply incremental numbers starting from 0, it would be more efficient use a list instead of a dict:

from itertools import product
trials = list(product(wgtlist, repeat=3))

or with an equivalent recursive generator:

def product(seq, repeat):
    if repeat:
        for i in seq:
            for p in product(seq, repeat - 1):
                yield [i, *p]
    else:
        yield []
blhsing
  • 91,368
  • 6
  • 71
  • 106
  • thank you! Is there a way to do this just using for loops and recursive function? or am I asking a stupid question? – localidiot Oct 27 '19 at 00:39
  • With `for` loops you won't be able to nest an arbitrary number of `for` loops. I've updated my answer with a recursive solution for your reference. – blhsing Oct 27 '19 at 00:54