AGC006C Rabbit Exercise(800) - tekiheiの日記

考察

うさぎ $j$ がジャンプする前後の座標の期待値の変化を考えます( $2 \le j \le N-1$ )。うさぎ $j$ のジャンプ前後のうさぎ $i(1 \le i \le N)$ の座標をそれぞれ $X_{i}, X_{i}'$ とします(これらは確率変数です)。
(1) $i=j$ のとき(うさぎ $i$ がジャンプする場合)
点 $x$ の点 $y$ に関して対象な点の座標を $f(x,y)$ で表すとすると、 \begin{eqnarray} E(X_{i}') &=& \displaystyle \sum_{x_{i}} \sum_{x_{i-1}} f(x_{i},x_{i-1}) \cdot P(X_{i}=x_{i},X_{i-1}=x_{i-1}) \cdot \frac{1}{2} + \sum_{x_{i}} \sum_{x_{i+1}} f(x_{i},x_{i+1}) \cdot P(X_{i}=x_{i},X_{i+1}=x_{i+1}) \cdot \frac{1}{2} \\ &=& \displaystyle \frac{1}{2} \sum_{x_{i}} \sum_{x_{i-1}} f(x_{i},x_{i-1}) \cdot P(X_{i}=x_{i},X_{i-1}=x_{i-1})+ \frac{1}{2} \sum_{x_{i}} \sum_{x_{i+1}} f(x_{i},x_{i+1}) \cdot P(X_{i}=x_{i},X_{i+1}=x_{i+1}) \\ &=& \displaystyle \frac{1}{2} E(f(x_{i},x_{i-1}))+ \frac{1}{2} E(f(x_{i},x_{i+1})) \\ &=& \frac{1}{2} E(2x_{i-1}-x_{i}) + \frac{1}{2} E(2x_{i+1}-x_{i}) \\ &=& \frac{1}{2}(2E(x_{i-1})-E(x_{i})) + \frac{1}{2}(2E(x_{i+1})-E(x_{i})) \\ &=& E(x_{i-1})+E(x_{i+1})-E(x_{i}) \end{eqnarray} 1行目は期待値の定義です(とりうる値 $\times$ 確率の総和)。2行目の $\frac{1}{2}$ の後ろはそれぞれ $f(x_{i},x_{i-1})$ と $f(x_{i},x_{i+1})$ の期待値の定義になっています。最後のほうは期待値の線形性を用いて変形しています。

(2) $i \neq j$ のとき
$X_{i}$ と $X_{i}'$ の確率分布は等しいので、期待値も等しいです。つまり $E(X_{i}')=E(X_{i})$ です(動かないので、それはそう?)。

(1)と(2)をまとめると、 $$ E(X_{i}')=\left\{ \begin{array}{ll} E(X_{i-1})+E(X_{i+1})-E(X_{i}) & (i=j) \\ E(X_{i}) & (i \neq j) \end{array} \right. $$

ジャンプする前のうさぎの位置の期待値はもとの座標に等しいので、結局以下の問題を解けば良いです。

$N$ 要素の数列 $(x_1, x_2, ... ,x_n)$ がある。 $j回目(1 \le j \le M)$ の操作では $i=a_{j}$ として $x_{i}$ を $x_{i-1}+x_{i+1}-x_{i}$ に置き換える。 $M$ 回の操作を $1$ セットとして $K$ セット行うとき、最終的な数列を求めよ。

操作を行列の積で表すことを考えます。例えば $N=5$ で $a_{i}=3$ だった場合、以下のように表せます。 $$ \begin{pmatrix} x_1' \\ x_2' \\ x_3' \\ x_4' \\ x_5' \\ \end{pmatrix}= \begin{pmatrix} 1 & 0 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 & 0 \\ 0 & 1 & -1 & 1 & 0 \\ 0 & 0 & 0 & 1 & 0 \\ 0 & 0 & 0 & 0 & 1 \\ \end{pmatrix} \begin{pmatrix} x_1 \\ x_2 \\ x_3 \\ x_4 \\ x_5 \\ \end{pmatrix} $$

よって、このような $M$ 個の行列の積を求めて $K$ 乗すれば、 $K$ セットの操作を表す行列が求まるので、最終的な数列を求められます。計算量は $O(N^{3}(M+\log K))$ です。これは少し改良することができます。上のような行列を左からかけることは「 $a_{i}行目を-1倍して、a_{i}-1行目とa_{i}+1行目をa_{i}行目に加える$ 」ことに相当します。よって $M$ 個の行列の積を求めるところを $O(N^{3}M)からO(NM)$ に改良することができます(TLEすることに変わりはないですが)。

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

using Lint=long long;

template<class T> class Matrix
{
 public:
    vector<vector<T>> A;
    Matrix(int n):A(n,vector<T>(n)){}
    Matrix(int n,int m):A(n,vector<T>(m)){}

    vector<T>& operator[](int i){ return A[i]; }
    int height(){ return A.size(); }
    int width(){ return A[0].size(); }
    static Matrix I(int n){
        Matrix mat(n);
        for(int i=0;i<n;i++) mat[i][i]=1;
        return mat;
    }

    Matrix& operator+=(Matrix& B){
        int H=height(), W=width();
        for(int i=0;i<H;i++) for(int j=0;j<W;j++)
            A[i][j]+=B[i][j];
        return *this;
    }
    Matrix& operator-=(Matrix& B){
        int H=height(), W=width();
        for(int i=0;i<H;i++) for(int j=0;j<W;j++)
            A[i][j]-=B[i][j];
        return *this;
    }
    Matrix& operator*=(Matrix& B){
        int H=height(), W=B.width(), K=width();
        vector<vector<T>> C(H,vector<T>(W));
        for(int i=0;i<H;i++) for(int j=0;j<W;j++)
            for(int k=0;k<K;k++) C[i][j]+=A[i][k]*B[k][j];
        A.swap(C);
        return *this;
    }
    Matrix operator+(Matrix& B){ return Matrix(*this)+=B; }
    Matrix operator-(Matrix& B){ return Matrix(*this)-=B; }
    Matrix operator*(Matrix& B){ return Matrix(*this)*=B; }
    Matrix pow(Lint K){
        Matrix res=Matrix::I(height());
        Matrix B=*this;
        while(K>0){
            if(K&1) res*=B; B*=B; K>>=1;
        }
        return res;
    }
};

int main()
{
   int N; cin>>N;
   vector<int> x(N);
   for(int i=0;i<N;i++) cin>>x[i];
   Lint M,K; cin>>M>>K;
   vector<int> a(M);
   for(int i=0;i<M;i++) cin>>a[i],a[i]--;

   assert((Lint)N*M+N*N*N<=(Lint)1e7); // これくらいまでは解ける?
   Matrix<Lint> A=Matrix<Lint>::I(N);
   for(int i=0;i<M;i++){
      for(int j=0;j<N;j++) A[a[i]][j]*=-1;
      for(int j=0;j<N;j++) A[a[i]][j]+=A[a[i]-1][j];
      for(int j=0;j<N;j++) A[a[i]][j]+=A[a[i]+1][j];
   }
   A=A.pow(K);

   Matrix<Lint> ans(N,1);
   for(int i=0;i<N;i++) ans[i][0]=x[i];
   ans=A*ans;
   for(int i=0;i<N;i++) cout<<ans[i][0]<<endl;
   return 0;
}

このままでは間に合わないので、別の方法を考えます。

「 $x_{i}$ を $x_{i}'=x_{i-1}+x_{i+1}-x_{i}$ に置き換える」という操作について考えてみます。 $x_{i-1}+x_{i+1}-x_{i}$ は $f(x_{i},x_{i-1})$ と $f(x_{i},x_{i+1})$ の中点です。よって図に書くと以下のようになります。

f:id:tekihei:20191225130254p:plain

ここで $x_{i-1}$ と $x_{i}'$ の距離はいくつになっているでしょうか？ $x_{i}'$ は中点なので、 $f(x_{i},x_{i-1})$ と $x_{i}'$ の距離は(赤)+(黄)です。よって $x_{i-1}$ と $x_{i}'$ の距離は(黄)になっていることが分かります。同様にして $x_{i}'$ と $x_{i+1}$ の距離が(赤)になっていることも分かります。

f:id:tekihei:20191225130301p:plain

以上よりこの操作は階差を交換する操作と言い換えることができそうです。実際 $d_{i} :=x_{i+1}-x_{i}$ とおくと
\begin{eqnarray} d_{i}' &=& x_{i+1}'-x_{i}' &=& x_{i+1}-(x_{i-1}+x_{i+1}-x_{i}) &=& x_{i}-x_{i-1} &=& d_{i-1} \\ d_{i-1}' &=& x_{i}'-x_{i-1}' &=& (x_{i-1}+x_{i+1}-x_{i})-x_{i-1} &=& x_{i+1}-x_{i} &=& d_{i} \end{eqnarray} より、(大小関係が図のようになっていない場合も)階差の交換と言い換えられることが分かります。

実際に階差の値を交換すると $O(MK)$ ですが、添字を交換することにすると、操作を $1$ から $N-1$ の置換で表すことができるので、置換の累乗で $O(M+N\log K)$ に高速化できます。あとは最終的な階差からもとの数列を復元することで答えを求められます。

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

using Lint=long long;

class Permutation
{
 private:
   vector<int> A;
 public:
   Permutation(int n):A(n){}
   int& operator[](int i){ return A[i]; }
   int size(){ return A.size(); }
   static Permutation I(int n){
      Permutation res(n);
      for(int i=0;i<n;i++) res[i]=i;
      return res;
   }

   Permutation& operator*=(Permutation& B){
      vector<int> C(size());
      for(int i=0;i<size();i++) C[i]=B[A[i]];
      A.swap(C);
      return *this;
   }
   Permutation operator*(Permutation& B){ return Permutation(*this)*=B; }
   Permutation pow(Lint K){
      Permutation res=Permutation::I(size());
      Permutation B=*this;
      while(K>0){ if(K&1) res*=B; B*=B; K>>=1; }
      return res;
   }
};

int main()
{
   int N; cin>>N;
   vector<int> x(N);
   for(int i=0;i<N;i++) cin>>x[i];
   int M; cin>>M;
   Lint K; cin>>K;
   vector<int> a(M);
   for(int i=0;i<M;i++) cin>>a[i],a[i]--;

   const int sz=N-1;
   vector<int> d(sz);
   for(int i=0;i<sz;i++) d[i]=x[i+1]-x[i];

   Permutation p=Permutation::I(sz);
   for(int i=0;i<M;i++){
      swap(p[a[i]-1],p[a[i]]);
   }
   p=p.pow(K);

   vector<Lint> ans(N);
   ans[0]=x[0];
   for(int i=0;i<sz;i++){
      ans[i+1]=ans[i]+d[p[i]];
   }
   for(int i=0;i<N;i++) cout<<ans[i]<<endl;
   return 0;
}

まとめ

期待値
行列の累乗は $O(N^{3} \log K)$ 、置換の累乗は $O(N \log K)$
図を書く

感想

結局対称な座標の中点にジャンプするとしてもよいみたいですが、式を書かないと理解できませんでした。期待値に確率をかけることに違和感を感じていましたが、期待値は $(とりうる値)\times(確率)$ の総和なので、この確率のほうに確率をかけていると考えるといい気がします。期待値と置換はまだあまり理解していないので他の問題も解こうと思います。