python - 三元表示中的快速数字总和(Python)

Question

我定义了一个函数

def enumerateSpin(n):
    s = []
    for a in range(0,3**n):
        ternary_rep = np.base_repr(a,3)
        k = len(ternary_rep)
        r = (n-k)*'0'+ternary_rep
        if sum(map(int,r)) == n:
            s.append(r)
    return s

我查看一个数字 0 <= a < 3^N 并询问其三进制表示形式中的数字之和是否达到某个值。我通过首先将数字转换为其三元表示形式的字符串来实现此目的。我填充零是因为我想存储一个固定长度表示的列表，稍后我可以将其用于进一步的计算(即两个元素之间的逐位比较)。

现在，np.base_repr 和 sum(map(int,#)) 在我的计算机上分别需要大约 5 us，这意味着一次迭代大约需要 10 us，并且我正在寻找一种方法，您可以完成我所做的事情，但速度要快 10 倍。

(编辑:关于在左侧填充零的注释)

(Edit2:事后看来，最终表示形式最好是整数元组而不是字符串)。

(Edit3:对于那些想知道的人来说，代码的目的是枚举具有相同总 S_z 值的 spin-1 链的状态。)

Answer 1

您可以使用itertools.product生成数字，然后转换为字符串表示形式:

import itertools as it

def new(n):
    s = []
    for digits in it.product((0, 1, 2), repeat=n):
        if sum(digits) == n:
            s.append(''.join(str(x) for x in digits))
    return s

这使我的速度提高了大约 7 倍:

In [8]: %timeit enumerateSpin(12)
2.39 s ± 7.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [9]: %timeit new(12)
347 ms ± 4.26 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

在 Python 3.9.0 (IPython 7.20.0) (Linux) 上测试。

上面的过程，使用it.product，也生成了一些数字，我们通过推理知道它们不遵守条件(这是所有数字的一半的情况，因为位数必须等于位数)。对于 n 位数字，我们可以计算 2、1 和 0 的各种数字计数，最终总计为n。然后我们就可以生成所有distinct permutations这些数字，因此只生成相关数字:

import itertools as it
from more_itertools import distinct_permutations

def new2(n):
    all_digits = (('2',)*i + ('1',)*(n-2*i) + ('0',)*i for i in range(n//2+1))
    all_digits = it.chain.from_iterable(distinct_permutations(d) for d in all_digits)
    return (''.join(digits) for digits in all_digits)

特别是对于大量的n，这会带来额外的、显着的加速:

In [44]: %timeit -r 1 -n 1 new(16)
31.4 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

In [45]: %timeit -r 1 -n 1 list(new2(16))
7.82 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

请注意，上述解决方案 new 和 new2 具有 O(1) 内存扩展(将 new 更改为 yield > 而不是追加)。