AtCoder Beginner Contest 172 D - Sum of Divisors

AtCoder 競技プログラミング約数

問題文

要約：
k=1からNまで、 $k * f(k)$ を足し上げる。
ただし、 $f(k)$ とは、正の整数であって、kの約数であるものの個数。

https://atcoder.jp/contests/abc172/tasks/abc172_d

O(N^2)

$1 \le j \le k$ なるすべての j について、ループを回して k の約数であるかを調べると $O(N^2)$ になります。

#include <iostream>
using ll = long long;

ll f_n2(int N)
{
    ll ans = 0;
    for (int k = 1; k <= N; ++k)
    {
        // 約数カウンタ
        int c = 0;
        for (int j = 1; j <= N; ++j)
        {
            // 割り切れるなら増やす
            if (k % j == 0) ++c;
        }
        ans += (ll)k * c;
    }
    return ans;
}

int main()
{
    int N;
    std::cin >> N;

    ll ans = f_n2(N);
    std::cout << ans << '\n';
}

N <= 10^7 ですので、当然、3secには間に合いません。

O(N√N)

約数を数える部分は $O(\sqrt N)$ にできます。

たとえば 20 の約数がいくつあるか数えることを考えましょう。
20 = 1 * 20
20 = 2 * 10
20 = 4 * 5

このように、割り切れたときには必ず相方がいることが分かります。
左側は必ず√20 より小さく、右側は必ず√20 より大きくなっているので、
(そうしないと、掛けたときに20になりません)
√20 = 4.47... まで探索して2倍すれば十分であることが分かります。

ただし、16のような平方数では、
16 = 1 * 16
16 = 2 * 8
16 = 4 * 4

このようになることがあるので、
√k ぴったりのときは1つしか数えないような工夫が必要です。

#include <iostream>
#include <cmath>

using ll = long long;

ll f_nsqn(int N)
{
    ll ans = 0;
    for (int k = 1; k <= N; ++k)
    {
        // 約数カウンタ
        int c = 0;
        int sqk = (int)sqrt(k);
        for (int j = 1; j <= sqk; ++j)
        {
            if (k % j == 0)
            {
                // k = j * j
                if (k / j == j)
                    c++;
                // それ以外
                else
                    c += 2;
            }
        }
        ans += (ll)k * c;
    }
    return ans;
}

int main()
{
    int N;
    std::cin >> N;

    ll ans = f_nsqn(N);
    std::cout << ans << '\n';
}

N = 10^7 のとき、N√N = 3 * 10^10 くらいになるので、これでも間に合いません。

O(NlogN)

k の約数を探しにいくのではなく、k 自身が将来の数の約数である、というふうに考えてみます。

たとえば、k=2 とすると、
2, 4, 6, 8, 10 はすべて 2 の倍数であり、
言い換えると、これらはすべて 2 を約数にもつ数でもあります。
したがって、f(2), f(4), f(6), f(8), f(10) は 2 が約数であることを知ることができます。
約数の数がひとつ増えたことになるので、+1 しておきます。
このようにすると、たとえば f(10) は、k=1, 2, 5, 10 のときに +1 されることになり、
これはつまり約数の数になっています。

図を見ると分かるように、これのループ回数は
N + N/2 + N/3 + ...... + N/N
になっています。

これの計算量を見積もるのは難しそうですが、
1/1 + 1/2 + 1/3 + ...... + 1/N
の形の和を調和級数と呼び、なんとその値はだいたいlogNになります。
したがって、この解法の計算量は $O(N\log N)$ です。

約数列挙の方針では「約数かどうかを判定してみたものの約数ではなかった」という場合があるのに対して、
こちらでは必ず約数になっています。直感的にも、こちらのほうが効率がよさそうです。

#include <iostream>
#include <vector>

using ll = long long;

ll f_nlogn_memo(int N)
{
    std::vector<int> p(N+1, 0);
    for (int k = 1; k <= N; ++k)
    {
        for (int j = k; j <= N; j += k)
        {
            // k, 2k, 3k, ... の約数カウントを増やす
            ++p[j];
        }
    }
    ll ans = 0;
    for (int k = 1; k <= N; ++k)
    {
        ans += (ll)k * p[k];
    }
    return ans;
}

int main()
{
    int N;
    std::cin >> N;
    ll ans = f_nlogn_memo(N);
    std::cout << ans << '\n';
}

C++でギリギリ通ります。
空間計算量がO(N)なので、10^7 ものメモリを確保するのが遅いのかもしれません。

O(NlogN) その2

もとの問題に立ち返ってみると、求めたいものは k * f(k) の総和でした。

ここまでの図をじーっと見ていると、
そもそも○を使う必要はなく、数字で埋めてしまっていいことがわかります。
この表中の数字をすべて足しさえすればよいので、
配列を使わなくても実装できます。

#include <iostream>
using ll = long long;

ll f_nlogn(int N)
{
    ll ans = 0;
    for (int k = 1; k <= N; ++k)
    {
        for (int j = k; j <= N; j += k)
        {
            // 数字を直接足し込む
            ans += j;
        }
    }
    return ans;
}

int main()
{
    int N;
    std::cin >> N;

    ll ans = f_nlogn(N);
    std::cout << ans << '\n';
}

C++では余裕をもって通ります。
言語によってはこれでもまだ通らないことがあるようです。
(Rubyで6300ms前後でした)

O(N)

「表中の数字をすべて足し上げる」という方針をもってさらに図を見つめていると、
横方向の足し算は等差数列の和であり、O(1) で計算できてしまうことが分かります。
横方向がO(1) で計算できたので、全体ではO(N)になります。

#include <iostream>
using ll = long long;

// 1からnまでの和
ll sn(ll n)
{
    return (ll)n * (n+1) / 2;
}

ll f_n(int N)
{
    ll ans = 0;
    for (int k = 1; k <= N; ++k)
    {
        ans += sn(N/k) * k;
    }
    return ans;
}

int main()
{
    int N;
    std::cin >> N;

    ll ans = f_n(N);
    std::cout << ans << '\n';
}

O(√N)

表中の数字をすべて足し上げればよいので、もはや見た目の順番通りに足す必要はありません。
表の数字を左に詰めると、対称性が見てとれます。
√N 個のL字型について、1つのL字は O(1) で計算できるので、
全体で O(√N) で計算ができます。

#include <iostream>
#include <cmath>

using ll = long long;

ll sn(ll n)
{
    return (ll)n * (n+1) / 2;
}

ll f_sqn(int N)
{
    int sqN = (int)sqrt(N);
    ll ans = 0;
    for (int k = 1; k <= sqN; ++k)
    {
        ans += (sn(N/k) - sn(k)) * k * 2 + k * k;
    }
    return ans;
}

int main()
{
    int N;
    std::cin >> N;

    ll ans = f_sqn(N);
    std::cout << ans << '\n';
}