test/yosupo/sqrt_mod.test.cpp

• View this file on GitHub
• Last update: 2020-11-13 20:22:13+09:00
• Problem: https://judge.yosupo.jp/problem/sqrt_mod

Depends on

• mod/sqrt.cpp

Code

// verification-helper: PROBLEM https://judge.yosupo.jp/problem/sqrt_mod

#include<bits/stdc++.h>
using namespace std;

#define call_from_test
#include "../../mod/sqrt.cpp"
#undef call_from_test

signed main(){
  cin.tie(0);
  ios::sync_with_stdio(0);

  int T;
  cin>>T;
  for(int t=0;t<T;t++){
    int y,p;
    cin>>y>>p;
    auto res=mod_sqrt(y,p);
    if(res.empty()) cout<<-1<<"\n";
    else cout<<res[0]<<"\n";
  }
  cout<<flush;
  return 0;
}

#line 1 "test/yosupo/sqrt_mod.test.cpp"
// verification-helper: PROBLEM https://judge.yosupo.jp/problem/sqrt_mod

#include<bits/stdc++.h>
using namespace std;

#define call_from_test
#line 1 "mod/sqrt.cpp"

#line 3 "mod/sqrt.cpp"
using namespace std;
#endif
//BEGIN CUT HERE
template<typename T>
int jacobi(T a,T m){
  int s=1;
  if(a<0) a=a%m+m;
  while(m>1){
    a%=m;
    if(a==0) return 0;
    int r=__builtin_ctz(a);
    if((r&1) and ((m+2)&4)) s=-s;
    a>>=r;
    if(a&m&2) s=-s;
    swap(a,m);
  }
  return s;
}

// MOD: prime
template<typename T>
vector<T> mod_sqrt(T a,const T MOD){
  if(MOD==2) return {a&1};
  int j=jacobi(a,MOD);
  if(j== 0) return {0};
  if(j==-1) return {};

  using ll = long long;
  mt19937 mt;
  ll b,d;
  while(1){
    b=mt()%MOD;
    d=(b*b-a)%MOD;
    if(d<0) d+=MOD;
    if(jacobi<ll>(d,MOD)==-1) break;
  }

  ll f0=b,f1=1,g0=1,g1=0;
  for(ll e=(MOD+1)>>1;e;e>>=1){
    if(e&1){
      ll tmp=(g0*f0+d*((g1*f1)%MOD))%MOD;
      g1=(g0*f1+g1*f0)%MOD;
      g0=tmp;
    }
    ll tmp=(f0*f0+d*((f1*f1)%MOD))%MOD;
    f1=(2*f0*f1)%MOD;
    f0=tmp;
  }
  if(g0>MOD-g0) g0=MOD-g0;
  return {T(g0),T(MOD-g0)};
}
//END CUT HERE
#ifndef call_from_test
//INSERT ABOVE HERE
signed main(){
  return 0;
}
#endif
#line 8 "test/yosupo/sqrt_mod.test.cpp"
#undef call_from_test

signed main(){
  cin.tie(0);
  ios::sync_with_stdio(0);

  int T;
  cin>>T;
  for(int t=0;t<T;t++){
    int y,p;
    cin>>y>>p;
    auto res=mod_sqrt(y,p);
    if(res.empty()) cout<<-1<<"\n";
    else cout<<res[0]<<"\n";
  }
  cout<<flush;
  return 0;
}

Back to top page