DIVISORSUM - Editorial

editorial
, ,
1

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: tle_99
Preparation: tabr
Tester: yash_daga
Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

Number theory

PROBLEM:

You’re given N and K.
Starting with A = [1, 2, \ldots, N], repeat the following process K times:

  • Replace each element of A with one copy of each of its divisors.

Find the sum of the final array.

EXPLANATION:

Notice that each initial element of A is independent: if we start with x, it’ll split into some divisors, then those divisors will further split, and so on.

Let f(x, K) denote the sum of the array formed if we start with the array [x] and perform the operation K times.
If we can compute all f(x, K) for 1 \leq x \leq N, the answer is their sum.

Note that f(x, K) can be defined recursively. Specifically,

f(x, K) = \begin{cases} x & \text{if } K = 0, \\ \sum_{d|N} f(d, K-1) & \text{otherwise} \end{cases}

This follows from the fact that x splits into one copy of each of its divisors; and then each of them independently goes through K-1 more steps.

The important observation to be made here is that f(x, K) is a multiplicative function in x.
That is, for two integers x and y such that \gcd(x, y) = 1, f(xy, K) = f(x, K)\times f(y, k).
This follows from the fact that if f is multiplicative, then so is \sum_{d\mid N} f(d). A proof can be found here.

This means it’s enough for us to compute f(x, K) for all those x that are prime powers; everything else follows.

Next, let’s figure out how to compute f(p^x, K) for a prime p.
As it turns out, that can be done recursively too!
In particular, we have:

f(p^x, K) = \begin{cases} 1 & \text{if } x = 0, \\ p\cdot f(p^{x-1}, K) + \binom{K-1+x}{K-1} & \text{otherwise} \end{cases}
Proof

We’ll prove this by induction on K.
For K = 0, recall that we have f(p^x, 0) = p^x.
But, p^x = p\cdot p^{x-1} + 0, and \binom{K-1+x}{K-1} is indeed 0 for K = 1 regardless of what x is.
So, the claim is true for K = 0.

Now, consider some K\gt 0.
Let’s apply the recursive formulation of f.

\begin{align*} f(p^x, K) &= \sum_{d \mid p^x} f(d, K-1) \\ &= \sum_{y=0}^x f(p^y, K-1) \\ &= 1 + \sum_{y=1}^x \left( p\cdot f(p^{y-1}, K-1) + \binom{K-2+y}{K-2} \right) \\ &= p\cdot\sum_{y=1}^x f(p^{y-1}, K-1) + \sum_{y=0}^x \binom{K-2+y}{K-2} \\ &= p\cdot f(p^{x-1}, K) + \binom{K-1+x}{K-1} \end{align*}

where the second summation is reduced via the hockey-stick identity.

This gives us rather simple-to-implement solution: for each prime p \leq N, compute the values for all its powers; then use these values to compute the values for all integers \leq N.
Both finding primes \leq N and propagating values from prime powers to other values can be done in a sieve-like fashion.

TIME COMPLEXITY:

\mathcal{O}(N\log\log N) per testcase.

CODE:

Editorialist's code (Python) 
mod = 10**9 + 7
inv = [1]*100
for i in range(1, 100): inv[i] = pow(i, mod-2, mod)

for _ in range(int(input())):
    n, k = map(int, input().split())
    dp = [-1]*(n+1)
    dp[1] = 1
    for p in range(2, n+1):
        if dp[p] != -1: continue
        x, pw = p, 1
        val = (p + k) % mod
        add = k
        while x <= n:
            for y in range(x, n+1, x):
                if (y//x)%p == 0: continue
                if dp[y] == -1: dp[y] = 1
                dp[y] = (dp[y] * val) % mod
            x *= p
            pw += 1
            add = (add * (k + pw - 1)) % mod
            add = (add * inv[pw]) % mod
            val = (val * p + add) % mod

    print(sum(dp[1:]) % mod)