Submission #601241

#	Time	Username	Problem	Language	Result	Execution time	Memory
601241		Petr	Star Trek (CEOI20_startrek)	C++17	100 / 100	200 ms	53304 KiB

startrek

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

const int mod=1e9+7;

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;
}

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 p, vector<vector<int> >& adi, vector<int>& vertex_status, vector<vector<int>>& red_neighbors){
    if (vertex_status[v] == -2) {
      vertex_status[v] = p;
      for(auto u: adi[v]){
          if(u==p) continue;
   
          if(!dfs_win(u, v, adi, vertex_status, red_neighbors)){
              red_neighbors[v].push_back(u);
          }
      }
    } else {
      if (vertex_status[v] >= 0 && vertex_status[v] != p) {
        auto u = vertex_status[v];
        if(!dfs_win(u, v, adi, vertex_status, red_neighbors)){
            red_neighbors[v].push_back(u);
        }
        vertex_status[v] = -1;
      }
    }
 
    if (vertex_status[v] == p) {
      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] != p;
    }
}
 
int dfs_redforce(int v, int p, vector<vector<int> >& adi, vector<vector<int>>& red_neighbors, vector<int>& vertex_status, vector<unordered_map<int, 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] != p;
 
    if(vertex_status[v]==-2){
        for(auto u: adi[v]){
            if(u==p) continue;
            redf[v][u]=dfs_redforce(u, v, adi, red_neighbors, vertex_status, redf, summe);
            summe[v]+=redf[v][u];
        }
 
        vertex_status[v]=p;
    }
    else{
        if(vertex_status[v]>=0 && vertex_status[v]!=p){
            auto u = vertex_status[v];
            redf[v][u]=dfs_redforce(u, v, adi, red_neighbors, vertex_status, redf, summe);
            summe[v]+=redf[v][u];
            vertex_status[v]=-1;
        }
    }
 
    int redforce = 0;
    if(!is_green){ //red
        redforce++;
        redforce+=summe[v];
        if (vertex_status[v] == -1 && p != -1) {
          redforce-=redf[v][p];
        }
    }
    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]==p) u=red_neighbors[v][1];
            else if(red_neighbors[v][1]!=p) needed=false;
        }
        if (needed) {
            redforce+=redf[v][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;
    vector<vector<int> > adi(n);
    for(int i=0; i<n-1; i++){
        int a,b;
        cin >> a >> b;
        adi[a-1].push_back(b-1);
        adi[b-1].push_back(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<unordered_map<int,int> > redf(n);
    vector<int> summe(n, 0);
    for(int i=n-1; i>=0; i--){
 
        bool we_win=dfs_win(i, -1, adi, vertex_status, red_neighbors);
        int redforce=dfs_redforce(i, -1, adi, 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:37: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]
   37 |   for (int i = 0; i < res.size(); ++i) res[i][i] = 1;
      |                   ~~^~~~~~~~~~~~
startrek.cpp: In function 'vect mul(matrix, vect)':
startrek.cpp:43: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]
   43 |   for (int i = 0; i < res.size(); ++i) {
      |                   ~~^~~~~~~~~~~~
startrek.cpp:44: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]
   44 |     for (int j = 0; j < res.size(); ++j) {
      |                     ~~^~~~~~~~~~~~
startrek.cpp: In function 'matrix multiply(matrix, matrix)':
startrek.cpp:53: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]
   53 |   for (int i = 0; i < res.size(); ++i) {
      |                   ~~^~~~~~~~~~~~
startrek.cpp:54: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]
   54 |     for (int j = 0; j < res.size(); ++j) {
      |                     ~~^~~~~~~~~~~~
startrek.cpp:55: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]
   55 |       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