Mr Tarsa
Mr Tarsa

Reputation: 6652

Generate all polynomial terms of certain degree

Having list of n terms

ts = ['t1','t2','t3',...,'tn']

there is a task to achieve all possible q-length combinations of this terms.

Thus, for

ts = ['t1','t2']
q = 4

the answer will be

[['t1','t1','t1','t1'],['t1','t2','t2','t2'],['t1','t1','t2','t2'],
 ['t1','t1','t1','t2'],['t2','t2','t2','t2']]

Here is my solution

powers = np.array(list(itertools.product(*[range(q+1)]*len(ts))))
powers = powers[np.where(np.asarray(map(sum, powers))==q)]
res = map(lambda ps_: flatten([ p*[t] for p,t in zip(ps_,ts) ]), powers)

where flatten is function that produces list from list of lists.

While having the solution that produces the desired result, I wonder if there are simpler ways to tackle the problem?

Upvotes: 1

Views: 528

Answers (1)

What you need are combinations with replacement. The simplest solution is to use the aptly named itertools.combinations_with_replacement:

>>> list(itertools.combinations_with_replacement(ts,q))
[('t1', 't1', 't1', 't1'),
 ('t1', 't1', 't1', 't2'),
 ('t1', 't1', 't2', 't2'),
 ('t1', 't2', 't2', 't2'),
 ('t2', 't2', 't2', 't2')]

Upvotes: 1

Related Questions