Submission #601251

#	Time	Username	Problem	Language	Result	Execution time	Memory
601251		Petr	Star Trek (CEOI20_startrek)	C++17	100 / 100	155 ms	30708 KiB

startrek

#include<bits/stdc++.h>
//#include<iostream>
//#include<vector>
using namespace std;
#define int long long

const int mod=1e9+7;

// Using a modint type instead of inserting "% mod" in appropriate
// places makes it easier to avoid both negative numbers and overflows.
struct modint {
  int val;

  modint() : val(0) {}

  modint(int val_) : val(val_ % mod) {
    if (val < 0) val += mod;
  }
}; 

modint operator+(modint a, modint b) {
  return a.val + b.val;
}

modint operator-(modint a, modint b) {
  return a.val - b.val;
}

modint operator*(modint a, modint b) {
  return a.val * b.val;
}

// Fixed-size arrays are nicer than vectors in that you don't
// need to resize(), and they also work much faster, so if we can
// use them, it's typically a good idea.
using vect = array<modint, 2>;
using matrix = array<vect, 2>;

matrix unit() {
  // Note that " = {};" is important to zero-initialize std::array.
  matrix res = {};
  for (int i = 0; i < res.size(); ++i) res[i][i] = 1;
  return res;
}

vect mul(matrix a, vect b) {
  vect res = {};
  for (int i = 0; i < res.size(); ++i) {
    for (int j = 0; j < res.size(); ++j) {
      res[i] = res[i] + a[i][j] * b[j];
    }
  }
  return res;
}

matrix multiply(matrix a, matrix b) {
  matrix res = {};
  for (int i = 0; i < res.size(); ++i) {
    for (int j = 0; j < res.size(); ++j) {
      for (int k = 0; k < res.size(); ++k) {
        res[i][j] = res[i][j] + a[i][k] * b[k][j];
      }
    }
  }
  return res;
}

matrix power(matrix a, int e){
    if(e==0){
        return unit();
    }
 
    matrix res=power(a, e/2);
    res=multiply(res, res);
    if(e%2==1){
        res=multiply(res, a);
    }
    return res;
}
 
 
bool dfs_win(int v, int peid, const vector<vector<int> >& adi, const vector<int>& dest, vector<int>& vertex_status, vector<vector<int>>& red_neighbors){
    if (vertex_status[v] == -2) {
      vertex_status[v] = peid;
      for(auto eid: adi[v]){
          if (eid == peid) continue;
          auto u = dest[eid];
          if(!dfs_win(u, eid ^ 1, adi, dest, vertex_status, red_neighbors)){
              red_neighbors[v].push_back(eid);
          }
      }
    } else {
      if (vertex_status[v] >= 0 && vertex_status[v] != peid) {
        auto eid = vertex_status[v];
        auto u = dest[eid];
        if(!dfs_win(u, eid ^ 1, adi, dest, vertex_status, red_neighbors)){
            red_neighbors[v].push_back(eid);
        }
        vertex_status[v] = -1;
      }
    }
 
    if (vertex_status[v] == peid) {
      return red_neighbors[v].size() > 0;
    } else {
      if (red_neighbors[v].size() >= 2) return true;
      if (red_neighbors[v].size() == 0) return false;
      return red_neighbors[v][0] != peid;
    }
}
 
int dfs_redforce(int v, int peid, const vector<vector<int> >& adi, const vector<int>& dest, const vector<vector<int>>& red_neighbors, vector<int>& vertex_status, vector<int>& redf, vector<int>& summe){
    bool is_green;
    if (red_neighbors[v].size() >= 2) is_green = true;
    else if (red_neighbors[v].size() == 0) is_green = false;
    else  is_green = red_neighbors[v][0] != peid;
 
    if(vertex_status[v]==-2){
        for(auto eid: adi[v]){
            if(eid==peid) {
              continue;
            }
            auto u = dest[eid];
            redf[eid]=dfs_redforce(u, eid ^ 1, adi, dest, red_neighbors, vertex_status, redf, summe);
            summe[v]+=redf[eid];
        }
 
        vertex_status[v]=peid;
    }
    else{
        if(vertex_status[v]>=0 && vertex_status[v]!=peid){
            auto eid = vertex_status[v];
            auto u = dest[eid];
            redf[eid]=dfs_redforce(u, eid ^ 1, adi, dest, red_neighbors, vertex_status, redf, summe);
            summe[v]+=redf[eid];
            vertex_status[v]=-1;
        }
    }
 
    int redforce = 0;
    if(!is_green){ //red
        redforce++;
        redforce+=summe[v];
        if (vertex_status[v] == -1 && peid != -1) {
          redforce-=redf[peid];
        }
    }
    else { //green with one red child
        bool needed=red_neighbors[v].size()<3;
        int u=red_neighbors[v][0];
        if(red_neighbors[v].size()==2){
            if(red_neighbors[v][0]==peid) u=red_neighbors[v][1];
            else if(red_neighbors[v][1]!=peid) needed=false;
        }
        if (needed) {
            redforce+=redf[u];
        }
    }
 
    return redforce;
 
}
 
signed main(){
    ios_base::sync_with_stdio(false);
    cin.tie(0);
 
    //freopen("input1.txt", "r", stdin);
    //freopen("out.txt", "w", stdout);
 
    int n, d;
    cin >> n >> d;
    // Now we store edge numbers here, and the target of the edge in "dest" array.
    // This allows us to avoid unordered_map. We also maintain the property
    // that the reverse to the edge x is the edge x^1.
    vector<vector<int> > adi(n);
    vector<int> dest(2 * (n - 1));
    for(int i=0; i<n-1; i++){
        int a,b;
        cin >> a >> b;
        adi[a-1].push_back(2 * i);
        dest[2 * i] = b - 1;
        adi[b-1].push_back(2 * i + 1);
        dest[2 * i + 1] = a - 1;
    }
 
    int G=0; int R=0;
    int g1=0; int r1=0;
    int c=0;
    vector<int> vertex_status(n, -2);
    vector<int> vertex_status2(n, -2);
    vector<vector<int>> red_neighbors(n);
    vector<int> redf(2 * (n - 1), 0);
    vector<int> summe(n, 0);
    for(int i=n-1; i>=0; i--){
 
        bool we_win=dfs_win(i, -1, adi, dest, vertex_status, red_neighbors);
        int redforce=dfs_redforce(i, -1, adi, dest, red_neighbors, vertex_status2, redf, summe);
 
        if(we_win){
            c++;
            g1=n;
            r1=n-redforce;
        }
        else{
            g1=0;
            r1=redforce;
        }
        G+=g1;
        R+=r1;
    }
 
    matrix m1 = {};
    m1[0][0] = g1;
    m1[0][1] = r1;
    m1[1][0] = n - g1; // This value is not important since we only print ans[0], but in case we also wanted ans[1] this would be correct.
    m1[1][1] = n - r1; // This value is not important since we only print ans[0], but in case we also wanted ans[1] this would be correct.
    matrix m2 = {};
    m2[0][0] = G;
    m2[0][1] = R;
    m2[1][0] = n * n - G;
    m2[1][1] = n * n - R;
    matrix p=power(m2, d-1);
    vect v=mul(p, {c, n-c});
    vect ans=mul(m1, v);
    cout << ans[0].val << "\n";
}

Compilation message (stderr)

startrek.cpp: In function 'matrix unit()':
startrek.cpp:42:21: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::array<std::array<modint, 2>, 2>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   42 |   for (int i = 0; i < res.size(); ++i) res[i][i] = 1;
      |                   ~~^~~~~~~~~~~~
startrek.cpp: In function 'vect mul(matrix, vect)':
startrek.cpp:48:21: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::array<modint, 2>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   48 |   for (int i = 0; i < res.size(); ++i) {
      |                   ~~^~~~~~~~~~~~
startrek.cpp:49:23: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::array<modint, 2>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   49 |     for (int j = 0; j < res.size(); ++j) {
      |                     ~~^~~~~~~~~~~~
startrek.cpp: In function 'matrix multiply(matrix, matrix)':
startrek.cpp:58:21: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::array<std::array<modint, 2>, 2>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   58 |   for (int i = 0; i < res.size(); ++i) {
      |                   ~~^~~~~~~~~~~~
startrek.cpp:59:23: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::array<std::array<modint, 2>, 2>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   59 |     for (int j = 0; j < res.size(); ++j) {
      |                     ~~^~~~~~~~~~~~
startrek.cpp:60:25: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::array<std::array<modint, 2>, 2>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   60 |       for (int k = 0; k < res.size(); ++k) {
      |                       ~~^~~~~~~~~~~~

• Subtask 1: 0 / 0
• Subtask 2: 7 / 7
• Subtask 3: 8 / 8
• Subtask 4: 15 / 15
• Subtask 5: 15 / 15
• Subtask 6: 20 / 20
• Subtask 7: 20 / 20
• Subtask 8: 15 / 15