SPCHEF - Editorial

iceknight1093 · August 30, 2023, 6:30pm

PROBLEM LINK:

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

Author: mikey_03
Tester & Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

Math, binary exponentiation

PROBLEM:

Chef wants to visit N temples, and offer an equal number of flowers to each one.
Before each temple, the number of flowers with him is multiplied by K.
After visiting all N temples, no flowers must remain.

Find the minimum number of flowers Chef must pluck so that this can be achieved, as well as the number of flowers to be offered at each temple.
Print the answer modulo M.

EXPLANATION:

We’ll do a bit of simple algebra to figure out the answer.

Suppose Chef picks x flowers initially, and offers y to each temple.
Let’s call F_i the number of flowers remaining with Chef after he’s visited the i-th temple.
Then, we know the following:

F_0 = x, i.e, Chef has x flowers to start with
For i \gt 0, F_i = K\cdot F_{i-1} - y, as per the statement.
F_N = 0, since no flowers must remain.

Looking at the values we get in this sequence,

F_0 = x \\ F_1 = xK - y \\ F_2 = K\cdot (xK - y) - y = xK^2 - yK - y \\ F_3 = K\cdot (xK^2 - yK - y) - y = xK^3 - yK^2 - yK - y \\ \vdots \\ F_i = xK^i - yK^{i-1} - yK^{i-2} - yK^{i-3} - \ldots - yK - y

Now, we also know that F_N = 0. Using that information, we see that

xK^N = y\cdot (1 + K + K^2 + \ldots + K^{N-1})

x and y must be positive integers, the hence the smallest value of x that satisfies this equation is (1 + K + K^2 + \ldots + K^{N-1}) itself; which in turn sets y = K^N.

Now that we know x and y, our task is to actually compute them modulo M quickly.

Subtask 1

The first subtask has prime M, which makes our task somewhat simple.
y = K^N can be computed in \mathcal{O}(\log N) using binary exponentiation, so let’s look at x.

We know that x = 1 + K + K^2 + \ldots + K^{N-1}
Notice that this is the sum of N terms of a geometric progression, for which there’s a well-known formula.
Applying that formula to our case, we see that x = \frac{K^N-1}{K-1}

The numerator is easy to compute: once again, K^N can be found in \mathcal{O}(\log N) using binary exponentiation.
Further, since M is prime (in particular, we need \gcd(M, K-1) = 1), “division” is also possible here, via the concept of modular inverses.

This allows us to “divide” K^N-1 by K-1, by computing the inverse of K-1 and multiplying with it instead.
Computing the inverse can be done in \mathcal{O}(\log{M}), as detailed in the above link.
So, subtask 1 is solved in \mathcal{O}(\log M) time.

Subtask 2

The second subtask has M be not necessarily a prime.
y = K^N can still be computed in exactly the same way, but our earlier method of computing x via inverses no longer works, because inverses may not exist for composite M.

However, with a bit of cleverness, we can still compute the sum of N terms of a geometric progression in \mathcal{O}(\log N).

Let f(a, N) = 1 + a + a^2 + \ldots + a^N.
Notice that the even and odd terms of the above sum themselves form a geometric progression, but with ratio a^2 instead.
That is, we have
f(a, N) = 1 + a + a^2 + \ldots + a^N = (1 + a^2 + a^4 + \ldots) + a\cdot(1 + a^2 + a^4 + \ldots)

In particular, working out odd N and even N separately, you can see that:

For odd N, f(a, N) = (1+a)\cdot f(a^2, \frac{N-1}{2})
For even N, f(a, N) = (1+a)\cdot f(a^2, \frac{N}{2}) - a^{N+1}
Alternately, you can say f(a, N) = 1 + a\cdot f(a, N-1) when N is even

Notice that we’re halving the exponent at each step, so along with the base case f(a, 0) = 1, we have an algorithm that takes \mathcal{O}(\log N) steps, which is exactly what we want!

TIME COMPLEXITY

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

CODE:

Author's code (C++, subtask 1)

#include<bits/stdc++.h>
 
using namespace std;
using ll = long long;
 
ll power(ll a,ll b,ll m){ 
  ll ans=1; while(b>0){ if(b&1) ans=(ans*a)%m; a=(a*a)%m; b>>=1; }return ans; 
} 
 
int inv(int a, int mod) {
  return a <= 1 ? a : mod - (long long)(mod/a) * inv((mod % a), mod) % mod;
}
   
void solve() 
{
    int n,k,mod; cin>>n>>k>>mod;
    
    int invKMinus1 = inv((k - 1), mod);
    
    int pluck = ((power(k,n,mod) - 1)%mod * (invKMinus1%mod)) % mod;
    
    int give = power(k,n,mod);
  
    cout<<pluck<<" "<<fixed<<setprecision(0)<<give<<'\n'; 
  
}
int main()
{
   ios_base::sync_with_stdio(0); cin.tie(NULL); cout.tie(NULL);
   int t;
   cin>>t;
   while(t--)
{
   solve();
}
}

Editorialist's code (C++)

#include "bits/stdc++.h"
using namespace std;
using ll = long long int;
mt19937_64 rng(chrono::high_resolution_clock::now().time_since_epoch().count());
 
int main()
{
    ios::sync_with_stdio(false); cin.tie(0);
 
    auto mpow = [&] (int a, int n, const int mod) {
        int res = 1;
        while (n) {
            if (n & 1) res = (1LL * res * a) % mod;
            a = (1LL * a * a) % mod;
            n /= 2;
        }
        return res;
    };
    auto gp_sum = [&] (const auto &self, int a, int n, const int mod) -> int {
        if (n == 0) return 1;
        int res = 1LL * (1 + a) * self(self, (1LL * a * a)%mod, n/2, mod) % mod;
        if (n%2 == 0) res = (res - mpow(a, n+1, mod) + mod) % mod;
        return res;
    };
    
    int t; cin >> t;
    while (t--) {
        int n, k, mod; cin >> n >> k >> mod;
        cout << gp_sum(gp_sum, k, n-1, mod) << ' ' << mpow(k, n, mod) << '\n';
    }
}

Editorialist's code (Python)

# Computes 1 + a + a^2 + ... + a^n, modulo mod
def gp_sum(a, n, mod):
    if n == 0: return 1
    if n%2 == 1: return (1 + a) * gp_sum(a*a % mod, n//2, mod) % mod
    else: return (1 + a * (a+1) % mod * gp_sum(a*a % mod, n//2 - 1, mod))%mod
 
for _ in range(int(input())):
    n, k, mod = map(int, input().split())
    print(gp_sum(k, n-1, mod), pow(k, n, mod))

nikok · August 30, 2023, 7:59pm

how do we know this is the smallest value?

roumak · August 31, 2023, 2:16am

I was able to come to the conclusion that the smallest number of flowers required would be a GP with a common ratio of k. But was unable to come up with a way to efficiently calculate the sum of GP (mod M). Because I found out that simply finding the sum and then using mod takes a lot of time and results in TLE.
Could you please expand on how the function used by the Editorialist is derived? or some link to articles where we can read about it in detail?

PS: The question was great!!

iceknight1093 · August 31, 2023, 4:29am

One way to see it is to note that you have

y = x\cdot\frac{K^N}{1 + K + K^2 + \ldots + K^{N-1}}

y has to be an integer, so the thing on the right also has to be an integer.
But, K^N and 1 + K + K^2 + \ldots + K^{N-1} don’t share any factors.

Why?

If K^N and 1 + K + K^2 + \ldots + K^{N-1} had a common factor, then they should also have a common prime factor, say p.
p is a prime factor of K^N, which means it’s also a prime factor of K itself.

But then, p can’t divide 1 + K + K^2 + \ldots + K^{N-1} (because it divides all the terms of that sum except the 1, and if it divided the full sum then it would have to divide 1 as well, which is impossible).
So, there can’t be a common prime factor, meaning there can’t be a common factor at all.

So, the only way for that quantity to be an integer, is if x is a multiple of (1 + K + \ldots + K^{N-1})
x needs to be positive, so of course the smallest multiple you can pick is (1 + K + \ldots + K^{N-1}) itself.

iceknight1093 · August 31, 2023, 4:31am

Which function do you mean?

roumak · August 31, 2023, 5:35am

iceknight1093:

def gp_sum(a, n, mod):
    if n == 0: return 1
    if n%2 == 1: return (1 + a) * gp_sum(a*a % mod, n//2, mod) % mod
    else: return (1 + a * (a+1) % mod * gp_sum(a*a % mod, n//2 - 1, mod))%mod

This function. I understand the objective of the function but also want to learn more about the math behind it.

iceknight1093 · August 31, 2023, 5:42am

That’s already explained in the editorial, in the “subtask 2” section.
What I called f(a, N) in the editorial is the gp_sum(a, n) function in my code.