PROBLEM LINK:
Author: Rahul Dugar
Tester: Ildar Gainullin
Editorialist: Rajarshi Basu
DIFFICULTY:
Medium-Hard
PREREQUISITES:
Discrete Math, Primitive Roots
PROBLEM:
You are given an integer A. Let’s define an infinite sequence S(A) = A\%P, A^2\%P, A^3\%P, \ldots, where P = 10^8+7 is a prime and \% is the modulo operator.
Let’s also define a decreasing sum D(S) as the sum of all elements of a sequence S which are strictly smaller than all preceding elements of S. When S is a sequence of non-negative integers, the number of such elements is clearly finite.
Find D(S(A)).
There are Q = 500 testcases.
EXPLANATION:
Some Initial Ideas?
If you know about discrete logarithm, the following sqrt-decomp-ish approach should come to naturally:
- First, we need a primitive root of P = 10^8 + 7, and then maybe we can see the elements of the sequence as powers of the primitive root. This is helpful since it allows us to represent the position as some discrete logarithm problem. 5 is a primitive root in this case.
- It seems very probable that as you generate the numbers of the sequence S, after \approx \sqrt{P} values, we will get a value that is less than \sqrt P.
- After we get this value, let’s say X \leq \sqrt P, we can generate the position of all the values 1,2,3 \dots X using the baby-step-giant-step algorithm.
Some Explanation
We know that a number X is present at the location k if A^k = X \mod P. However, we do not need to do this computation for every query. We can precompute. Let’s say the primitive root is g. Then calculate y where g^y = X \mod P for all X \leq \sqrt P. Now, for our query let’s say that A = g^z. So we want to find A^k = X, or (g^z)^k = g^y which means kz = y, which means k = yz^{-1}. If we just find z^{-1} once for every query, we can calculate k in O(1) for every number.
Note that, finding inverse of z might not be so easy here since we are calculating inverse w.r.t P-1 here. Hence, we need to do a few additional steps. Let’s say that f = gcd(z,P-1). For k to exist, then f | y. If it doesn’t, that means that k does not exist. Let’s assume f | y. Then divide all z,P-1, and y by f. Now, we can find \frac{z}{f}^{-1} \mod \frac{P-1}{f}, and this serves as our “ z^{-1} ” that we were using earlier to multiply. Look at setter’s code for the exact implementation.
Once we find the position of the numbers from 1,2,…X, and the terms which appear before X, we can find the required sum. Think why this is enough!
Hint: After X, the next number has to be less than X, and it has to keep decreasing. So 1…X are the only interesting numbers.
- If we just use the normal baby-step-giant-step, the complexity will be O(\sqrt P \log(p) + Q*\sqrt P + \sqrt P * \sqrt P *log(P)).
- The first term is from calculation of factors for giant step,
- the second term is for bruting the first \approx \sqrt P naively in hopes that we get a term less than \sqrt P.
- The third term is to precalculate the discrete \log of the first \sqrt P values.
- Overall, it’s clear that this time complexity is not enough to pass the full solution.
Is everything above even actually correct
The main assumption in the above solution is that we will reach something lesser than \sqrt P in \approx \sqrt P steps. Even though it’s hard to prove, one can just run some codes and notice that the maximum number of steps needed to reach a number around \leq 10^5 is \approx 19,000, which seems pretty decent. So the idea is correct, but if only we could make it faster \dots
The crucial idea
Let us write the complexity in a bit more elaborate fashion. Recall that in the baby -step-giant-step algorithm, when we write x = np-q, we choose n = \sqrt P , hence giving us the \sqrt P terms. Instead, let us write the complexity without setting n to any particular value. Then, our complexity becomes the following:
- O(\frac{p}{n} \log P + Q*\sqrt P + \sqrt P * n * \log P).
- Note that the terms are exactly in the order as given in the section "Some Initial Ideas”
Can you find a better value for n?
Click for answer
Set n = \sqrt[4] P and see the magic !!
For implementation details, look at setter’s code.
SOLUTION:
Setter’s Code
#include <bits/stdc++.h>
using namespace std;
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/rope>
using namespace __gnu_pbds;
using namespace __gnu_cxx;
#ifndef rd
#define trace(...)
#define endl '\n'
#endif
#define pb push_back
#define fi first
#define se second
#define int long long
typedef long long ll;
typedef long double f80;
#define double long double
#define pii pair<int,int>
#define pll pair<ll,ll>
#define sz(x) ((long long)x.size())
#define fr(a,b,c) for(int a=b; a<=c; a++)
#define rep(a,b,c) for(int a=b; a<c; a++)
#define trav(a,x) for(auto &a:x)
#define all(con) con.begin(),con.end()
const ll infl=0x3f3f3f3f3f3f3f3fLL;
const int infi=0x3f3f3f3f;
const int mod=998244353;
//const int mod=1000000007;
typedef vector<int> vi;
typedef vector<ll> vl;
typedef tree<pii, null_type, less<pii>, rb_tree_tag, tree_order_statistics_node_update> oset;
auto clk=clock();
mt19937_64 rang(chrono::high_resolution_clock::now().time_since_epoch().count());
int rng(int lim) {
uniform_int_distribution<int> uid(0,lim-1);
return uid(rang);
}
int powm(int a, int b, int mod) {
int res=1;
while(b) {
if(b&1)
res=(res*a)%mod;
a=(a*a)%mod;
b>>=1;
}
return res;
}
vector<pii> gs;
vi facs;
const int p=100000007,root=5;
bool check(int a) {
for(int i:facs)
if(powm(a,(p-1)/i,p)==1)
return 0;
return 1;
}
int drt(int a) {
for(int i=0; i<1e8;i++) {
if((*lower_bound(all(gs),pii{a,0})).fi==a)
return ((*lower_bound(all(gs),pii{a,0})).se-i+p-1)%(p-1);
a=(a*root)%p;
}
}
int gcdExtended(int a, int b, int *x, int *y) {
if(a==0) {
*x=0,*y=1;
return b;
}
int x1,y1;
int gcd=gcdExtended(b%a,a,&x1,&y1);
*x=y1-(b/a)*x1;
*y=x1;
return gcd;
}
int modInverse(int a, int m) {
int x,y;
int g=gcdExtended(a,m,&x,&y);
return (x%m+m)%m;
}
int firstfew[25005];
void solve() {
int small=pow(p,0.25);
int big=p/small;
int lol=1;
for(int i=0; i<small; i++)
lol=(lol*root)%p;
int poo=lol;
for(int i=1; i<big; i++) {
gs.pb({poo,i*small});
poo=(poo*lol)%p;
}
sort(all(gs));
int sqt=sqrt(p);
for(int i=1; i<=sqt; i++)
firstfew[i]=drt(i);
int q;
cin>>q;
while(q--) {
int a;
cin>>a;
int ans=a+1;
int last=a,best=a;
while(last>sqt) {
last=(last*a)%p;
if(last<best) {
best=last;
ans+=best;
}
}
if(last==1)
ans--;
int rt=drt(a);
int common=__gcd(p-1,rt);
int modhere=(p-1)/common;
int inv=modInverse(rt/common,modhere);
int min_encountered=p;
for(int i=2; i<best; i++) {
int lol=firstfew[i];
if(lol%common)
continue;
lol=((lol*inv)/common)%modhere;
if(lol<min_encountered) {
ans+=i;
min_encountered=lol;
}
}
cout<<ans<<'\n';
}
}
signed main() {
ios_base::sync_with_stdio(0),cin.tie(0),cout.tie(0);
srand(chrono::high_resolution_clock::now().time_since_epoch().count());
cout<<fixed<<setprecision(8);
int t=1;
// cin>>t;
fr(i,1,t)
solve();
#ifdef rd
cerr<<endl<<endl<<endl<<endl<<"Time elapsed: "<<(double)(clock()-clk)/CLOCKS_PER_SEC<<endl;
#endif
return 0;
}