Yukicoder No.278 連続する整数の和(2)

solution

約数の総和を求める問題に帰着。$O(\sqrt{N} \log \log N)$。

$a_0 + (a_0 + 1) + (a_0 + 2) + \dots + (a_0 + N - 1) = Na_0 + \frac{N(N+1)}{2}$である。 $a_0$は勝手に動くため$X$には影響できず、問題の$X$は$X \mid \mathrm{gcd}(N, \frac{N(N+1)}{2})$であるような$X$である。 このような$X$の総和とは、$\mathrm{gcd}(N, \frac{N(N+1)}{2})$の約数の総和。 $\mathrm{gcd}(N, \frac{N(N+1)}{2})$は$N$の偶奇により$N$または$\frac{N}{2}$である。約数の総和は素因数分解を用いて求めればよい。

implementation

#include <iostream>
#include <vector>
#include <map>
#include <cmath>
#define repeat(i,n) for (int i = 0; (i) < (n); ++(i))
typedef long long ll;
using namespace std;
vector<int> sieve_of_eratosthenes(int n) { // enumerate primes in [2,n] with O(n log log n)
    vector<bool> is_prime(n+1, true);
    is_prime[0] = is_prime[1] = false;
    for (int i = 2; i*i <= n; ++i)
        if (is_prime[i])
            for (int k = i+i; k <= n; k += i)
                is_prime[k] = false;
    vector<int> primes;
    for (int i = 2; i <= n; ++i)
        if (is_prime[i])
            primes.push_back(i);
    return primes;
}
map<ll,int> factors(ll n, vector<int> const & primes) {
    map<ll,int> result;
    for (int p : primes) {
        if (n < p *(ll) p) break;
        while (n % p == 0) {
            result[p] += 1;
            n /= p;
        }
    }
    if (n != 1) result[n] += 1;
    return result;
}
ll f(ll x) {
    vector<int> primes = sieve_of_eratosthenes(sqrt(x) + 3);
    ll result = 1;
    for (auto it : factors(x, primes)) {
        ll p; int cnt; tie(p, cnt) = it;
        ll acc = 1;
        ll e = 1; repeat (i, cnt) { e *= p; acc += e; }
        result *= acc;
    }
    return result;
}
int main() {
    ll n; cin >> n;
    ll x = n % 2 == 0 ? n/2 : n;
    cout << f(x) << endl;
    return 0;
}