CODEWITHSUNDEEP
pythonlistlist-comprehension
Leo
Leo

Reputation: 288

Python nesting list comprehensions with variable depth

I am writing a program where I use different methods to fit a dataset, and in the final step I want to take a distribution over the models, and then test it against a validation set to pick the optimal distribution. In order to do so, I need lists that sum up to 1 (the total weight of all the models). In the case of 3 models, I use the following code:

Grid = np.arange(0,1.1,0.1)
Dists = [[i,j,k] for i in Grid for j in Grid for k in Grid if i+j+k==1]

I am now looking for a way to generalize this to arbitrary number of models, say d, without specifying what d is beforehand. I have looked at np.tensordot and np.outer, but couldnt figure out a way to make this work. Any ideas would be appreciated. cheers, Leo

Upvotes: 1

Views: 156

Answers (1)

Ashwini Chaudhary
Ashwini Chaudhary

Reputation: 250921

You are looking for itertools.product:

from itertools import product
Dists = [list(p) for p in product(Grid, repeat=3) if sum(p) == 1]

Upvotes: 2

Related Questions