Quiz
https://atcoder.jp/contests/agc028/tasks/agc028_b
Nが小さい場合
- 解説放送の通り期待値を考える。最初はdoubleを使って確率を定義通りに求めてみる
- サンプルテストケースは全て通る。提出すると当然TLE
vector<ll> A(N); FOR(i, 0, N){ cin >> A.at(i); } double expectation = 0; FOR(i, 0, N){ // a_i が出現する回数の期待値 double p = 0; FOR(j, 0, N){ ll n = abs(j-i)+1; p += 1.0 / n; } expectation += A[i] * p; } ll ans = round(expectation * factorial(N)); ans %= mod; p(ans);
Nが大きい場合
- 解説pdfにある通り、mod上で全て行う
- mod上では割り算を扱うときはフェルマーの小定理が必要
- ACしたコード https://atcoder.jp/contests/agc028/submissions/4473940
- doubleは全て消え、llになる
- 割り算は逆元をかけることになるので小数点は出てこない
- pを求める部分を逆元で書くと、下記コードのコメントアウトされた箇所になる
- これでは全体でO(N2)なので、累積和でO(N)にしたコードが下記である
Code
#include<algorithm> #include<complex> #include<ctype.h> #include<iomanip> #include<iostream> #include<map> #include<math.h> #include<numeric> #include<queue> #include<set> #include<stack> #include<stdio.h> #include<string> #include<string> #include<vector> using namespace std; typedef long long ll; #define FOR(i,a,b) for(ll i=(a);i<(b);++i) #define ALL(v) (v).begin(), (v).end() #define p(s) cout<<(s)<<endl #define p2(s, t) cout << (s) << " " << (t) << endl #define br() p("") #define pn(s) cout << (#s) << " " << (s) << endl #define p_yes() p("YES") #define p_no() p("NO") template < typename T > void vprint(T &V){ for(auto v : V){ cout << v << " "; } cout << endl; } const ll mod = 1e9 + 7; const ll inf = 1e18; ll factorial(ll n){ ll ret = 1; FOR(i, 1, n+1){ ret *= i; } return ret; } const int N_MAX = 100010; ll Per[N_MAX] = {}; // n! ll Inv[N_MAX] = {}; // 1/2, 1/3, ... ll Ac[N_MAX] = {}; // Invの累積和 // a^b mod p ll mod_pow(ll a, ll b){ if(b==0) return 1; // 肩が奇数 if(b%2==1){ return a * mod_pow(a, b-1) % mod; } else{ return mod_pow(a*a % mod, b/2) % mod; } } int main(){ cin.tie(0); ios::sync_with_stdio(false); // input ll N; cin >> N; vector<ll> A(N); FOR(i, 0, N){ cin >> A.at(i); } // nCr高速化準備 Per[1] = 1; FOR(i, 2, N_MAX){ Per[i] = i * Per[i-1] % mod; } Inv[1] = 1; FOR(i, 2, N_MAX){ Inv[i] = mod_pow(i, 1000000005); } // 累積和 Ac[1] = Inv[1]; FOR(i, 2, N_MAX){ Ac[i] = Ac[i-1] + Inv[i]; } ll expectation = 0; FOR(i, 0, N){ ll p = 0; // FOR(j, 0, N){ // ll n = abs(j-i)+1; // p += Inv[n]; // p %= mod; // } p += Ac[N-1 - i + 1]; // 右 p %= mod; p += Ac[i-0 + 1] - Ac[1]; // 左 p %= mod; expectation += A[i] * p; expectation %= mod; } ll ans = expectation * Per[N]; ans %= mod; p(ans); return 0; }
解くまでに必要なステップ
- N!回ではなく1回あたりの期待値を求めればいいと気づく
- 期待値を正しく求める
- 解が大きすぎるのでmodした結果を答えるが、割り算が出てくるのでフェルマーの小定理を使う
- 累積和でオーダーを落とす