CF 1970 E3

题目描述

有 \(N\) 个小屋，每条小屋有 \(s_i\) 条短径，\(l_i\) 条长径通向湖。

每天，你将从一个小屋出发，先走到湖再走到另一个小屋（可以是一开始的小屋），但其中必须至少有一条通过的小径是短径。

一开始你在 \(1\) 号小屋，求走 \(M\) 天的方案数。

思路

由于把小屋当做状态的时间复杂度至少 \(O(N^2)\)，所以考虑从湖开始走。

我们令 \(dp_{i,0/1}\) 表示在第 \(i\) 天回到湖，走的最后一条道路是短/长径的方案数。

我们有转移：

\[\begin{array}{l} dp_{i,0}=dp_{i-1,0}\cdot\sum \limits_{j=1}^N (s_j+l_j)\cdot s_j+dp_{i-1,1}\cdot\sum \limits_{j=1}^N s_j^2\\ dp_{i,1}=dp_{i-1,0}\cdot\sum \limits_{j=1}^N (s_j+l_j)\cdot l_j+dp_{i-1,1}\cdot\sum \limits_{j=1}^N s_j\cdot l_j \end{array} \]

这个转移可以用矩阵乘法加速，这样每次乘法只有 \(2^2=O(1)\)。

空间复杂度 \(O(N)\)，时间复杂度 \(O(N+\log M)\)。

代码

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

const int MAXN = 100001, MOD = int(1e9) + 7;

struct Matrix {
  int n, m, mat[3][3];
  void Clear(int a, int b) {
    n = a, m = b;
    for(int i = 1; i <= n; ++i) {
      for(int j = 1; j <= m; ++j) {
        mat[i][j] = 0;
      }
    }
  }
  Matrix operator*(const Matrix &b) {
    Matrix ret;
    ret.Clear(n, b.m);
    for(int i = 1; i <= n; ++i) {
      for(int j = 1; j <= b.m; ++j) {
        for(int k = 1; k <= m; ++k) {
          ret.mat[i][j] = (ret.mat[i][j] + 1ll * mat[i][k] * b.mat[k][j] % MOD) % MOD;
        }
      }
    }
    return ret;
  }
}mat;

int n, m, a[MAXN], b[MAXN], ans;

Matrix Pow(Matrix x, Matrix a, int b) {
  for(; b; a = a * a, b >>= 1) {
    if(b & 1) {
      x = x * a;
    }
  }
  return x;
}

int main() {
  ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
  cin >> n >> m;
  for(int i = 1; i <= n; ++i) {
    cin >> a[i];
  }
  for(int i = 1; i <= n; ++i) {
    cin >> b[i];
  }
  mat.Clear(2, 2);
  for(int i = 1; i <= n; ++i) {
    mat.mat[1][1] = (mat.mat[1][1] + 1ll * (a[i] + b[i]) * a[i] % MOD) % MOD;
    mat.mat[1][2] = (mat.mat[1][2] + 1ll * (a[i] + b[i]) * b[i] % MOD) % MOD;
    mat.mat[2][1] = (mat.mat[2][1] + 1ll * a[i] * a[i] % MOD) % MOD;
    mat.mat[2][2] = (mat.mat[2][2] + 1ll * a[i] * b[i] % MOD) % MOD;
  }
  Matrix res;
  res.Clear(1, 2);
  res.mat[1][1] = a[1], res.mat[1][2] = b[1];
  Matrix x = Pow(res, mat, m - 1);
  for(int i = 1; i <= n; ++i) {
    ans = (0ll + ans + 1ll * x.mat[1][1] * (a[i] + b[i]) % MOD + 1ll * x.mat[1][2] * a[i] % MOD) % MOD;
  }
  cout << ans;
  return 0;
}