记一次筛素数算法的优化

求出1..n中有多少个素数的算法我们通常用筛数法，一种比较简单的实现如下：

#include <stdio.h>
#include <stdlib.h>

int main(const int argc, const char *argv[], const char *envp[])
{
    int64_t n;
    int64_t i, j;
    int64_t *arr;
    int64_t cnt = 0;

    if (argc != 2) {
        return -1;
    }

    sscanf(argv[1], "%lld", &n);

    arr = (int64_t *)calloc(n + 1, sizeof(int64_t));

    for (i = 2; i <= n; i++) {
        arr[i] = 1;
    }

    for (i = 2; i < n; i++) {    // outer for-loop

        if (arr[i] == 0) {
            continue;
        }

        for (j = i + i; j <= n; j += i) {    // inner for-loop
            arr[j] = 0;
        }
    }

    for (i = 2; i <= n; i++) {
        cnt += arr[i] != 0;
    }

    printf("%lld
", cnt);

    return 0;
}

这里首先解释一下为什么外层循环不用取n，因为如果n是合数，显然n已经在i为n最小素因数的时候就已经被筛掉了；否则n是素数不用筛。

取n=100000000测测运行时间：

 ~/ time ./test1 100000000
5761455
./test1 100000000  2.92s user 0.23s system 99% cpu 3.157 total

现在我们注意到：每次进入内层for循环的时候，实际上对于素数i，所有比i²小的i的倍数{j * i | j < i}都一定有一个比i更小的素因数，那我们其实可以将j从i²开始迭代，这样就可以节约一定的时间，简易实现如下：

#include <stdio.h>
#include <stdlib.h>

int main(const int argc, const char *argv[], const char *envp[])
{
    int64_t n;
    int64_t i, j;
    int64_t *arr;
    int64_t cnt = 0;

    if (argc != 2) {
        return -1;
    }

    sscanf(argv[1], "%lld", &n);

    arr = (int64_t *)calloc(n + 1, sizeof(int64_t));

    for (i = 2; i <= n; i++) {
        arr[i] = 1;
    }

    for (i = 2; i < n; i++) {    // outer for-loop

        if (arr[i] == 0) {
            continue;
        }

        for (j = i * i; j <= n; j += i) {    // inner for-loop
            arr[j] = 0;
        }
    }

    for (i = 2; i <= n; i++) {
        cnt += arr[i] != 0;
    }

    printf("%lld
", cnt);

    return 0;
}

仅仅改动了一个字符（'+' → '*'），有多大的时间收益呢？

 ~/ time ./test2 100000000
5761455
./test2 100000000  2.02s user 0.25s system 99% cpu 2.263 total

时间减少了近1s，几乎节约了1/3的时间，这个收效很明显啊。而且从输出来看，确实应该是正确的。（我知道单个测试样例不能证明程序的正确性，但是对于筛素数的算法而言，n取得足够大还能对，这就是一种佐证）

那我们再来思考思考还能不能再有所提升？

我们继续考虑内层循环，其对arr[j]的访问实际上相当没有空间连续性，并且因此也失掉了向量化的可能性。考虑到进入内层循环时所有比i小的素数的倍数都已经被筛过了，所以对于i的在[i, n/i]范围内的倍数中的也有很大部分被筛过了，也就是说我们仅需要筛掉[i, n/i]中尚未被筛掉的数的i倍的那个数{j * i | j ∈ [i, n/i]，且j未被筛掉}。这样的话我们首先可以得出一个推论：外层循环只需要取到sqrt(n)即可，实际中我直接用i * i <= n来作为外层循环的条件。另外，我们应当尽量消除条件分支，由于判断j是否被筛掉会造成一个if语句，所以我们可以改成内存循环中始终将arr[k]置0，若j未被筛掉则k = j * i，否则k = 0，综合起来就是k = arr[j] * j * i，这样就去掉了if语句并且利用了arr[0]留空这点，同时还使得如果不需要筛掉j * i的情况有很好的空间局部性。更进一步的，由于总是会计算arr[j] * j，所以我们的arr[]不要仅用来存放{0, 1}，而是直接用arr[j]存放j，这样k就是arr[j] * i。这样程序就变成了：

#include <stdio.h>
#include <stdlib.h>

int main(const int argc, const char *argv[], const char *envp[])
{
    int64_t n;
    int64_t i, j;
    int64_t *arr;
    int64_t cnt = 0L;

    if (argc != 2) {
        return -1;
    }

    sscanf(argv[1], "%lld", &n);

    arr = (int64_t *)calloc(n + 1, sizeof(int64_t));

    for (i = 2; i <= n; i++) {
        arr[i] = i;
    }

    for (i = 2; i * i <= n; i++) {    // outer for-loop

        if (arr[i] == 0) {
            continue;
        }

        for (j = n / i; j >= i; j--) {    // inner for-loop
            const int64_t k = arr[j] * i;
            arr[k] = 0;
        }
    }

    for (i = 2; i <= n; i++) {
        cnt += arr[i] != 0;
    }

    printf("%lld
", cnt);

    return 0;
}

之所以这里内存循环变成了降序的迭代，是因为避免j和j * i同时都没被筛掉过的话，且j * i² < n的话，会引起错误。

这次程序改动较大，那有没有较大的时间上的收益呢？

 ~/ time ./test3 100000000
5761455
./test3 100000000  1.01s user 0.25s system 99% cpu 1.263 total

又减少了1s，只不过这1s优化来得比较麻烦。