在Python中迭代使用常量求和的索引

时间:2016-11-21 20:13:59

标签: python loops iteration itertools

我正在使用多项式并且必须执行一些操作,其中我的循环变量的总和低于某个常量d。 现在我有

for i in range(0, d):
    for j in range(i, d):
        for k in range(j, d):

这对我来说似乎有点难看。 是否有一些函数,可能在itertools中,允许我迭代for i, j, k in foo(d)

2 个答案:

答案 0 :(得分:3)

你可以自己写。这是3个变量的强力方式:

def constant_sum(s):
    for i in range(s+1):
        for j in range(s-i+1):
            k = s - i - j
            yield i,j,k

def inferior_sum(s):
    for i in range(s+1):
        for j in range(s+1):
            if i+j >= s:
                break
            for k in range(s+1):
                if i+j+k < s:
                    yield i,j,k
                else:
                    break

for i,j,k in constant_sum(3):
    print(i,j,k)
print()
for i,j,k in inferior_sum(3):
    print(i,j,k)

输出:

0 0 3
0 1 2
0 2 1
0 3 0
1 0 2
1 1 1
1 2 0
2 0 1
2 1 0
3 0 0

0 0 0
0 0 1
0 0 2
0 1 0
0 1 1
0 2 0
1 0 0
1 0 1
1 1 0
2 0 0

以下是递归版本,可以为任何总和做任意数量的变量(n)...轻度测试:

def constant_sum(n,s):
    if n == 1:
        yield [s]
    else:
        for i in range(s+1):
            for r in constant_sum(n-1,s-i):
                yield [i] + r

def inferior_sum(n,s):
    if n == 1:
        for i in range(s):
            yield [i]
    else:
        for i in range(s):
            for r in inferior_sum(n-1,s-i):
                yield [i] + r

for x in constant_sum(3,3):
    print(*x)
print()
for x in inferior_sum(3,3):
    print(*x)

输出:

0 0 3
0 1 2
0 2 1
0 3 0
1 0 2
1 1 1
1 2 0
2 0 1
2 1 0
3 0 0

0 0 0
0 0 1
0 0 2
0 1 0
0 1 1
0 2 0
1 0 0
1 0 1
1 1 0
2 0 0

答案 1 :(得分:2)

itertools / functional方式类似于:

from itertools import product
inferior_sum3 = filter(lambda x: sum(x)<3, product(range(4),range(4),range(4)))
for permu in inferior_sum3:
    print(permu)

输出:

(0, 0, 0)
(0, 0, 1)
(0, 0, 2)
(0, 1, 0)
(0, 1, 1)
(0, 2, 0)
(1, 0, 0)
(1, 0, 1)
(1, 1, 0)
(2, 0, 0)