CodeIQの七夕問題を解いてみた

CodeIQの七夕問題で提出したコードを貼っておく。

問題と解説はこちら。
給料UPを目指す彦星の解答は？「七夕問題☆牽牛 彦星 牛をもっと飼う」解説 #CodeIQ

言語は C。
64ビットの整数型 long long int と立方根 cbrt() が使える処理系が必要。
実行時間は Ideone.com で 0.01秒程度。

(こっちに の場合のコードを置いておいた。実行時間は 0.9秒程度)

#include <stdio.h>
#include <math.h>
#include <string.h>
#include <time.h>

/* 牛舎のサイズ */
#define N 7777777

/* 以下の P_MAX, P_TABLE_SIZE, HASH_SIZE は N に応じて手動で適当な値に調節する必要がある */

/*
    P_MAX 以下の素数について素数表を作る
    P_MAX ≧ 3√N であれば動くが、この P_MAX の値が実行速度に影響するので実験して決めた
*/
#define P_MAX (N / 4)

/* 素数表のサイズ。P_MAX以下の素数の数 + 10 以上であればよい */
#define P_TABLE_SIZE 150000

/* 2の冪乗でないと動かない。N=10^7 で 0x10000、N=10^9 で 0x100000 あれば十分 */
#define HASH_SIZE 0x10000

#define SIEVE_SIZE (P_MAX / 30 + 1)
#define square(x) ((long long) (x) * (x))

int pi(int n);
int prime[P_TABLE_SIZE] = {2, 3, 5}, n_primes;
struct {int key; int value;} hash_table[HASH_SIZE];

/* 素数表 (prime[]) を作る。まだ高速化の余地あり */
int make_prime_table(void)
{
    static char sieve[SIEVE_SIZE];
    const int b[] = {1, 7, 11, 13, 17, 19, 23, 29},
    mask[] = {1, 0, 0, 2, 0, 4, 8, 0, 16, 32, 0, 64, 0, 0, 128};
    int i, j, m, p, size = 1, s, t, k = 0, n = 3, lim = sqrt(P_MAX);

    while (1) {
        for (k++; sieve[k >> 3] & 1 << (k & 7); k++)
            ;
        if ((p = 30 * (k >> 3) + b[k & 7]) > lim) break;
        if (size < SIEVE_SIZE) {
            s = p * size > SIEVE_SIZE ? SIEVE_SIZE : p * size;
            while (size < s) {
                t = 2 * size > s ? s : 2 * size;
                strncpy(sieve + size, sieve, t - size);
                size = t;
            }
        }
        for (j = 0; j < 8; j++) {
            i = p / 30 * p + b[j] * p / 30;
            if (i < p * p / 30) i += p;
            m = mask[b[j] * p % 30 / 2];
            for ( ; i < size; i += p) sieve[i] |= m;
        }
        prime[n++] = p;
        sieve[p / 30] |= mask[p % 30 / 2];
    }

    sieve[0] |= 1;
    for (i = k >> 3; i < SIEVE_SIZE; i++)
        for (j = 0, m = ~sieve[i] & 0xff; m; j++, m >>= 1)
            if (m & 1) prime[n++] = 30 * i + b[j];

    while (prime[n - 1] > P_MAX) n--;
    return n;
}

int hash_search(int key)
{
    /* ハッシュ関数を知らないわけではない */
    int idx = key & (HASH_SIZE - 1);
    while (hash_table[idx].key && hash_table[idx].key != key) {
        if (idx)
            idx--;
        else
            idx = HASH_SIZE - 1;
    }
    return idx;
}

/* prime[k] ≦ n ＜ prime[k+1] となる k を素数表から求めて返す */
int prime_index(int n)
{
    int a = 0, b = n_primes, c;
    while (b - a > 1) {
        c = (a + b) / 2;
        if (prime[c] > n)
            b = c;
        else
            a = c;
    }
    return a;
}

/*
    p, q, r, … を prime[k], prime[k+1], prime[k+2], …, prime[lim] として
    n/p + n/q + n/r + … - n/p/q - n/p/r - … - n/q/r - … + n/p/q/r + …
    を返す
*/
int g(int n, int k, int lim)
{
    int j, m, mk, u, mu, c, mc, sqrt_n, p = prime[k], sum = 0;

    if (n < p) return 0;

    if (n < prime[lim]) {
        u = pi(n);
        mu = 1;
    } else {
        u = lim;
        mu = n / prime[lim];
    }

    /*
        Σ_j n/prime[j] の計算
        j は k から u まで動かす
        mc を適当に決めて
        n/prime[j] ＞ mc の j については上の和をそのまま計算 (1)
        n/prime[j] ≦ mc の j については m = n/prime[j] として、同じ m のものをまとめて計算 (2)
        (1), (2) に該当する j が存在しないことがあるので、場合分けが必要
    */
    sqrt_n = sqrt(n);
    mc = sqrt_n / 3;    /* この式は勘と試行錯誤で決めた */
    mk = n / p;
    if (mc < mk) {
        c = pi(n / (mc + 1));
        if (c > u) c = u;    /* (2) に該当する j が存在しない */
        /* (1) の計算 */
        for (j = k; j <= c; j++)
            sum += n / prime[j];
    } else {    /* (1) に該当する j が存在しない */
        c = k - 1;
        mc = mk;
    }

    /* (2) の計算 */
    if (c < u) {
        sum += u * mu - c * mc;
        for (m = mu + 1; m <= mc; m++)
            sum += pi(n / m);
    }

    /* 2個以上の素数が分母にくる項の計算 */
    for (j = k; prime[j] <= sqrt_n; j++)
        sum -= g(n / prime[j], j + 1, lim);

    return sum;
}

/*
    k2 = pi(sqrt(n)), k3 = pi(cbrt(n)) として呼び出す
    n 以下の自然数のうち、prime[k3] ＜ p, q ≦ prime[k2] の
    ふたつの素数 p, q の積 pq の形の合成数の個数
*/
int h(int n, int k2, int k3)
{
    int k, sum = (k2 - k3) * (1 -  k2 - k3) / 2;
    for (k = k3 + 1; k <= k2; k++)
        sum += pi(n / prime[k]);
    return sum;
}

/*
    n 以下の素数の数 (π(n) と書くあれ) から 1 引いたもの
    ハッシュにはこの関数の値を記録しておく
*/
int pi(int n)
{
    int pi_n, k3, hash_idx = hash_search(n);
    if (hash_table[hash_idx].key) return hash_table[hash_idx].value;
    if (n <= P_MAX) {    /* n が小さいときは素数表から計算 */
        pi_n = prime_index(n);
    } else {
        k3 = pi(cbrt(n));
        pi_n = n + k3 - 1 - g(n, 0, k3) - h(n, pi(sqrt(n)), k3);
    }
    /* pi を再帰的に呼んでいるので、低確率で計算済みの値を上書きすることがあるが、効率が落ちるだけでバグりはしない */
    hash_table[hash_idx].key = n;
    return hash_table[hash_idx].value = pi_n;
}

/*
    p, q, r, … を prime[k], prime[k+1], prime[k+2], … として
    (n/p)^2 + (n/q)^2 + (n/r)^2 + … - (n/p/q)^2 - (n/p/r)^2 - … - (n/q/r)^2 - … + (n/p/q/r)^2 + …
    を返す
    計算の仕方、変数の意味は g とほとんど同じ。素数の上限がないぶん単純
*/
long long f(int n, int k)
{
    int j, m, mk, c, mc, sqrt_n, p = prime[k];
    long long sum = 0;

    if (n < p) return sum;

    sqrt_n = sqrt(n);
    mc = sqrt_n / 3;
    mk = n / p;
    if (mc < mk) {
        c = pi(n / (mc + 1));
        for (j = k; j <= c; j++)
            sum += square(n / prime[j]);
    } else {
        c = k - 1;
        mc = mk;
    }

    for (m = 1; m <= mc; m++)
        sum += (2 * m - 1) * (pi(n / m) - c);

    for (j = k; prime[j] <= sqrt_n; j++)
        sum -= f(n / prime[j], j + 1);

    return sum;
}

int main(void)
{
    time_t start = clock();
    n_primes = make_prime_table();
    printf("%lld\n", 2 + square(N - 1) - f(N - 1, 0));
    printf("elapsed time = %f sec\n", (float) (clock() - start) / CLOCKS_PER_SEC);
    return 0;
}

このコードを実行すると答である 36775823261333 が得られる。
因みに牛舎のサイズ N が大きくなると、飼える牛の頭数/ は に収束する。理由は、例えば、Wikipedia の「互いに素」の項参照。

だから、上の結果でほぼ小数点以下6位まで一致している。