加快 Python 中的位串/位操作？

Question

我使用 埃拉托色尼筛法 和 Python 3.1 编写了一个素数生成器。该代码在 ideone.com 上以 0.32 秒的速度正确而优雅地运行，以生成高达 1,000,000 的素数。

 # from bitstring import BitString

def prime_numbers(limit=1000000):
    '''Prime number generator. Yields the series
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...
    using Sieve of Eratosthenes.
    '''
    yield 2
    sub_limit = int(limit**0.5)
    flags = [False, False] + [True] * (limit - 2)
#     flags = BitString(limit)
    # Step through all the odd numbers
    for i in range(3, limit, 2):
        if flags[i] is False:
#        if flags[i] is True:
            continue
        yield i
        # Exclude further multiples of the current prime number
        if i <= sub_limit:
            for j in range(i*3, limit, i<<1):
                flags[j] = False
#                flags[j] = True

问题是，当我尝试生成高达 1,000,000,000 的数字时，内存不足。

     flags = [False, False] + [True] * (limit - 2)
MemoryError

可以想象，分配 10 亿个布尔值（在 Python 中每个1 字节4 或 8 字节（见注释））确实不可行，所以我研究了 bitstring 。我认为，为每个标志使用 1 位会更节省内存。然而，该程序的性能急剧下降——运行时间为 24 秒，对于高达 1,000,000 的素数。这可能是由于位串的内部实现。

您可以注释/取消注释这三行以查看我更改为使用 BitString 的内容，如上面的代码片段。

我的问题是， 有没有办法加快我的程序，有或没有位串？

编辑：请在发布前自行测试代码。我自然不能接受运行速度比我现有代码慢的答案。

再次编辑：

我已经在我的机器上编制了一份基准测试列表。

原文由 Xavier Ho 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

您的版本有几个小的优化。通过颠倒 True 和 False 的角色，您可以将“ if flags[i] is False: ”更改为“ if flags[i]: ”。第二个 range 语句的起始值可以是 i*i 而不是 i*3 。您的原始版本在我的系统上需要 0.166 秒。通过这些更改，以下版本在我的系统上需要 0.156 秒。

 def prime_numbers(limit=1000000):
    '''Prime number generator. Yields the series
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...
    using Sieve of Eratosthenes.
    '''
    yield 2
    sub_limit = int(limit**0.5)
    flags = [True, True] + [False] * (limit - 2)
    # Step through all the odd numbers
    for i in range(3, limit, 2):
        if flags[i]:
            continue
        yield i
        # Exclude further multiples of the current prime number
        if i <= sub_limit:
            for j in range(i*i, limit, i<<1):
                flags[j] = True

但是，这对您的记忆问题没有帮助。

进入 C 扩展的世界，我使用了 gmpy 的开发版本。 （免责声明：我是维护者之一。）开发版本称为 gmpy2，支持称为 xmpz 的可变整数。使用 gmpy2 和以下代码，我的运行时间为 0.140 秒。 1,000,000,000 限制的运行时间为 158 秒。

 import gmpy2

def prime_numbers(limit=1000000):
    '''Prime number generator. Yields the series
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...
    using Sieve of Eratosthenes.
    '''
    yield 2
    sub_limit = int(limit**0.5)
    # Actual number is 2*bit_position + 1.
    oddnums = gmpy2.xmpz(1)
    current = 0
    while True:
        current += 1
        current = oddnums.bit_scan0(current)
        prime = 2 * current + 1
        if prime > limit:
            break
        yield prime
        # Exclude further multiples of the current prime number
        if prime <= sub_limit:
            for j in range(2*current*(current+1), limit>>1, prime):
                oddnums.bit_set(j)

推动优化并牺牲清晰度，我使用以下代码获得 0.107 和 123 秒的运行时间：

 import gmpy2

def prime_numbers(limit=1000000):
    '''Prime number generator. Yields the series
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...
    using Sieve of Eratosthenes.
    '''
    yield 2
    sub_limit = int(limit**0.5)
    # Actual number is 2*bit_position + 1.
    oddnums = gmpy2.xmpz(1)
    f_set = oddnums.bit_set
    f_scan0 = oddnums.bit_scan0
    current = 0
    while True:
        current += 1
        current = f_scan0(current)
        prime = 2 * current + 1
        if prime > limit:
            break
        yield prime
        # Exclude further multiples of the current prime number
        if prime <= sub_limit:
            list(map(f_set,range(2*current*(current+1), limit>>1, prime)))

编辑：基于这个练习，我修改了 gmpy2 以接受 xmpz.bit_set(iterator) 。使用以下代码，对于小于 1,000,000,000 的所有素数，Python 2.7 的运行时间为 56 秒，Python 3.2 的运行时间为 74 秒。 （如评论中所述， xrange 比 range 快。）

 import gmpy2

try:
    range = xrange
except NameError:
    pass

def prime_numbers(limit=1000000):
    '''Prime number generator. Yields the series
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...
    using Sieve of Eratosthenes.
    '''
    yield 2
    sub_limit = int(limit**0.5)
    oddnums = gmpy2.xmpz(1)
    f_scan0 = oddnums.bit_scan0
    current = 0
    while True:
        current += 1
        current = f_scan0(current)
        prime = 2 * current + 1
        if prime > limit:
            break
        yield prime
        if prime <= sub_limit:
            oddnums.bit_set(iter(range(2*current*(current+1), limit>>1, prime)))

编辑#2：再试一次！我修改了 gmpy2 以接受 xmpz.bit_set(slice) 。使用以下代码，对于 Python 2.7 和 Python 3.2，所有小于 1,000,000,000 的素数的运行时间约为 40 秒。

 from __future__ import print_function
import time
import gmpy2

def prime_numbers(limit=1000000):
    '''Prime number generator. Yields the series
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...
    using Sieve of Eratosthenes.
    '''
    yield 2
    sub_limit = int(limit**0.5)
    flags = gmpy2.xmpz(1)
    # pre-allocate the total length
    flags.bit_set((limit>>1)+1)
    f_scan0 = flags.bit_scan0
    current = 0
    while True:
        current += 1
        current = f_scan0(current)
        prime = 2 * current + 1
        if prime > limit:
            break
        yield prime
        if prime <= sub_limit:
            flags.bit_set(slice(2*current*(current+1), limit>>1, prime))

start = time.time()
result = list(prime_numbers(1000000000))
print(time.time() - start)

编辑 #3：我更新了 gmpy2 以正确支持 xmpz 位级别的切片。性能没有变化，但 API 非常好。我做了一些调整，我把时间减少到大约 37 秒。 （请参阅编辑 #4 以更改 gmpy2 2.0.0b1。）

 from __future__ import print_function
import time
import gmpy2

def prime_numbers(limit=1000000):
    '''Prime number generator. Yields the series
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...
    using Sieve of Eratosthenes.
    '''
    sub_limit = int(limit**0.5)
    flags = gmpy2.xmpz(1)
    flags[(limit>>1)+1] = True
    f_scan0 = flags.bit_scan0
    current = 0
    prime = 2
    while prime <= sub_limit:
        yield prime
        current += 1
        current = f_scan0(current)
        prime = 2 * current + 1
        flags[2*current*(current+1):limit>>1:prime] = True
    while prime <= limit:
        yield prime
        current += 1
        current = f_scan0(current)
        prime = 2 * current + 1

start = time.time()
result = list(prime_numbers(1000000000))
print(time.time() - start)

编辑 #4：我在 gmpy2 2.0.0b1 中做了一些改变，打破了前面的例子。 gmpy2 不再将 True 视为提供无限源 1 位的特殊值。应该改用 -1。

 from __future__ import print_function
import time
import gmpy2

def prime_numbers(limit=1000000):
    '''Prime number generator. Yields the series
    2, 3, 5, 7, 11, 13, 17, 19, 23, 29 ...
    using Sieve of Eratosthenes.
    '''
    sub_limit = int(limit**0.5)
    flags = gmpy2.xmpz(1)
    flags[(limit>>1)+1] = 1
    f_scan0 = flags.bit_scan0
    current = 0
    prime = 2
    while prime <= sub_limit:
        yield prime
        current += 1
        current = f_scan0(current)
        prime = 2 * current + 1
        flags[2*current*(current+1):limit>>1:prime] = -1
    while prime <= limit:
        yield prime
        current += 1
        current = f_scan0(current)
        prime = 2 * current + 1

start = time.time()
result = list(prime_numbers(1000000000))
print(time.time() - start)

编辑 #5：我对 gmpy2 2.0.0b2 做了一些改进。您现在可以迭代所有设置或清除的位。运行时间提高了约 30%。

 from __future__ import print_function
import time
import gmpy2

def sieve(limit=1000000):
    '''Returns a generator that yields the prime numbers up to limit.'''

    # Increment by 1 to account for the fact that slices do not include
    # the last index value but we do want to include the last value for
    # calculating a list of primes.
    sieve_limit = gmpy2.isqrt(limit) + 1
    limit += 1

    # Mark bit positions 0 and 1 as not prime.
    bitmap = gmpy2.xmpz(3)

    # Process 2 separately. This allows us to use p+p for the step size
    # when sieving the remaining primes.
    bitmap[4 : limit : 2] = -1

    # Sieve the remaining primes.
    for p in bitmap.iter_clear(3, sieve_limit):
        bitmap[p*p : limit : p+p] = -1

    return bitmap.iter_clear(2, limit)

if __name__ == "__main__":
    start = time.time()
    result = list(sieve(1000000000))
    print(time.time() - start)
    print(len(result))

原文由 casevh 发布，翻译遵循 CC BY-SA 3.0 许可协议

Answer 2

好吧，这是我的第二个答案，但由于速度至关重要，我想我不得不提到 bitarray 模块——尽管它是 bitstring 的克星 :)。它非常适合这种情况，因为它不仅是 C 扩展（并且比纯 Python 有希望的速度更快），而且它还支持切片分配。

我什至没有尝试优化它，我只是重写了位串版本。在我的机器上，对于一百万以下的素数，我得到 0.16 秒。

对于十亿，它运行得非常好，并在 2 分 31 秒内完成。

 import bitarray

def prime_bitarray(limit=1000000):
    yield 2
    flags = bitarray.bitarray(limit)
    flags.setall(False)
    sub_limit = int(limit**0.5)
    for i in range(3, limit, 2):
        if not flags[i]:
            yield i
            if i <= sub_limit:
                flags[3*i:limit:i*2] = True

原文由 Scott Griffiths 发布，翻译遵循 CC BY-SA 4.0 许可协议