小さい数n（100万までとか）の素因数分解

CodeChefのフォーラムで June Challenge 2017の問題 Chef and Prime Queries (PRMQ) の解法を読んでいて
discuss.codechef.com

素因数分解をするのに

This can be done by creating a Smallest Prime Factor array in the sieve function itself.

とあって。エラトステネスの篩 (Sieve of Eratosthenes) をするときに、篩と同時に一番小さい素因数のテーブルを作ってしまえば素因数分解が簡単、っていう話。
篩の配列を一番小さい素因数を表すテーブルにしてしまえばいいよね。
（なんで今までそうしてこなかったんだ…今まで愚直に、2から までの素数でnを割り切れるかいちいち１つずつ試してた。これなら100倍ぐらい速そう）

#include <vector>
using namespace std;

vector<int> primes;
vector<int> smallest_prime_factor;

int sieve(int nmax){
    primes.clear();
    smallest_prime_factor.assign(nmax+1, 0);

    // ここではあえて愚直な書き方
    for (int n=2; n<=nmax; ++n) {
        if (!smallest_prime_factor[n]) {
            primes.push_back(n);
            for (int kn=n; kn<=nmax; kn+=n) {
                if (!smallest_prime_factor[kn]) {  // ←このチェックをせずに上書きを続けると largest_prime_factor[] が出来上がる
                    smallest_prime_factor[kn] = n;
                }
            }
        }
    }
    return primes.size();
}

// xを素因数分解
vector<int> factorize(int x) {
    vector<int> prime_factors;
    while (x > 1) {
        int s = smallest_prime_factor[x];
        if (!s) break;
        prime_factors.push_back(s);
        x /= s;
    }
    return prime_factors;
}

例えば x=120 の場合、

smallest_prime_factor[120] = 2
120 / 2 = 60
smallest_prime_factor[60] = 2
60 / 2 = 30
smallest_prime_factor[30] = 2
30 / 2 = 15
smallest_prime_factor[15] = 3
15 / 3 = 5
smallest_prime_factor[5] = 5
5 / 5 = 1

なので、factorize(120) は vector<int> {2,2,2,3,5} を返す。

nが100万までなら prime_factors[] の長さは高々19 で、
使われている素数は高々7個

本番中に書いたPRMQのコードの素因数分解部分をこれに差し替えてみたら、N=100000のケースで10〜20msecほど縮まった。
でもlargeでTLEを出してた（segment tree導入前の）コードは、1.3秒ほどかかる手元の最悪ケースで10msec縮んでもどうしようもない。

おまけ

in scheme

(define (sieve nmax)
  (let ((smallest-prime-factor (make-vector (+ 1 nmax) #f)))
    (let loop ((n 2) (primes ()))
      (if (< nmax n) (values (reverse! primes) smallest-prime-factor)
          (if (vector-ref smallest-prime-factor n)
              (loop (+ n 1) primes)
              (let next ((kn n))
                (if (< nmax kn) (loop (+ n 1) (cons n primes))
                    (begin
                      (if (vector-ref smallest-prime-factor kn) #f
                          (vector-set! smallest-prime-factor kn n))
                      (next (+ kn n))))))))))

(define (factorize smallest-prime-factor x)
  (let loop ((x x) (prime-factors ()))
    (if (= x 1) (reverse! prime-factors)
        (let ((s (vector-ref smallest-prime-factor x)))
          (loop (/ x s) (cons s prime-factors))))))

#|
(receive (primes smallest-prime-factor) (sieve 100000)
  (print primes)
  (print smallest-prime-factor)
  (print (factorize smallest-prime-factor 12345))  ; -> (3 5 823)
  )
|#

追記

最初に書いたschemeコードは (vector-ref smallest-prime-factor kn) をチェックしていなかったせいで、（smallest-じゃなくて）largest-prime-factorを生成してた。
いやそれでもうまく行く（し、そっちの方が微妙に速い）んだけど。
factorizeが逆順になる（のでreverse!もしなくていいし更に微妙に速いか）。

(define (sieve nmax)
  (let ((largest-prime-factor (make-vector (+ 1 nmax) #f)))
    (let loop ((n 2) (primes ()))
      (if (< nmax n) (values (reverse! primes) largest-prime-factor)
          (if (vector-ref largest-prime-factor n)
              (loop (+ n 1) primes)
              (let next ((kn n))
                (if (< nmax kn) (loop (+ n 1) (cons n primes))
                    (begin
                      (vector-set! largest-prime-factor kn n)
                      (next (+ kn n))))))))))

(define (factorize largest-prime-factor x)
  (let loop ((x x) (prime-factors ()))
    (if (= x 1) prime-factors
        (let ((s (vector-ref largest-prime-factor x)))
          (loop (/ x s) (cons s prime-factors))))))

#|
(receive (primes largest-prime-factor) (sieve 100000)
  (print primes)
  (print largest-prime-factor)
  (print (factorize largest-prime-factor 12345))  ; -> (3 5 823)
  )
|#