三元表示法中的快速数字和（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：事后看来，最终表示形式最好是整数的元组，而不是字符串）

（编辑3：对于那些想知道的人来说，代码的目的是枚举具有相同总S_z值的自旋-1链的状态。）

Answer 1

通常，要获取特定基数中的数字，我们可以执行以下操作：

while num > 0:
    digit = num % base
    num //= base
    print(digit)

使用num = 14, base = 3运行此命令时，我们会得到：

2
1
1

这意味着三元中的14是112。
我们可以将其提取到方法digits(num, base)中，并且仅在实际需要将数字转换为字符串时使用np.base_repr(a,3)：

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

---

为enumerateSpin(4)输出：

['0022', '0112', '0121', '0202', '0211', '0220', '1012', '1021', '1102', '1111', '1120', '1201', '1210', '2002', '2011', '2020', '2101', '2110', '2200']