Submission #607805

#	Time	Username	Problem	Language	Result	Execution time	Memory
607805		tmt514	Wells (CEOI21_wells)	C++14	30 / 100	10018 ms	59048 KiB

wells

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

const int MAXN = 1500001;
const long long MOD = 1e9+7;
vector<int> a[MAXN];
int e[MAXN][2];
int dp[MAXN][2];

int isfree[MAXN];
int depth[MAXN];
int N, K;

// compute longest path on edge i direction d
int go(int i, int d) {
  if(dp[i][d] != -1) return dp[i][d];
  // fprintf(stderr, "i=%d, d=%d\n", i, d);
  int x = e[i][d], y = e[i][d^1];
  dp[i][d] = 2;
  for(int j : a[y]) {
    if (j != i) {
      int d2 = (y == e[j][1]);
      dp[i][d] = max(dp[i][d], go(j, d2) + 1);
    }
  }
  return dp[i][d];
}

// test
bool test(int x, int p, int level, int longest_outside_path=0) {
  
  bool success = true;
  if (level == 0) {
    for(int i : a[x]) {
      int d = (x == e[i][1]);
      int y = e[i][d^1];
      if (y == p) continue;
      if (isfree[y]) continue;

      success &= test(y, x, K-1, dp[i][d^1]);
      if (!success) return false;
    }
  } else {
    int deepmost = 0, deeptwo = 0, deepthree = 0;
    for(int i : a[x]) {
      int d = (x == e[i][1]);
      int y = e[i][d^1];
      if (y == p) continue;
      if (isfree[y]) continue;

      success &= test(y, x, level - 1, dp[i][d^1]);
      if (!success) return false;

      if (dp[i][d]-1 > deepmost) {
        deepthree = deeptwo;
        deeptwo = deepmost;
        deepmost = dp[i][d] - 1;
      } else if (dp[i][d]-1 > deeptwo) {
        deepthree = deeptwo;
        deeptwo = dp[i][d] - 1;
      } else if (dp[i][d] - 1 > deepthree) {
        deepthree = dp[i][d] - 1;
      }
    }
    if (deepmost >= level && deeptwo >= level && deepmost + deeptwo + 1 >= K && level * 2 + 1 <= K) {
      return false;
    }
    
    if (deeptwo + longest_outside_path >= K  &&
      min(level-1, deepmost) + min(level-1, deeptwo) + 1 >= K) {
      // printf("x=%d, deep=%d + %d\n", x, min(level-1,deepmost), min(level-1,deeptwo));
      return false;
    }
  }
  return success;
}

int main() {
  scanf("%d%d", &N, &K);
  for(int i=1;i<N;i++) {
    int x, y;
    scanf("%d%d", &x, &y);
    a[x].push_back(i);
    a[y].push_back(i);
    e[i][0] = x;
    e[i][1] = y;
  }
  memset(dp, -1, sizeof(dp));
  for(int i=1;i<N;i++) {
    for(int d=0;d<2;d++) go(i, d);
  }

  
  // check if each vertex is a "free vertex"  
  int freecount=0;
  for(int x=1;x<=N;x++) {
    vector<int> lens;
    for(int i : a[x]) {
      int d = (x == e[i][1]);
      lens.push_back(dp[i][d]-1);
    }
    sort(lens.begin(), lens.end(), greater<>());
    int big = (lens.size()>=1? lens[0]:0) + (lens.size()>=2? lens[1]:0);
    if(big+1<K) {
      isfree[x]=1;
      ++freecount;
    }
  }

  // trivial case: no vertex belongs to a path of length K.
  if (freecount == N) {
    long long ans = 1;
    for(int i=0;i<N;i++) ans=ans*2%MOD;
    printf("YES\n%lld\n", ans);
    return 0;
  }
  // find one non-free vertex, and do DFS.
  // find a path of length k.
  auto get_one_path_start = [&]() {
    for(int x=1;x<=N;x++) {
    for(int i: a[x]) {
      int d = (x==e[i][1]);
      if(dp[i][d] >= K)
      return x;
    }  
  }
  return -1;
  };
  int x = get_one_path_start();
  vector<int> path;
  while(path.size() < K-1) {
    for(int i:a[x]) {
      int d = (x==e[i][1]);
      if((path.size()==0 || e[i][d^1]!=path.back()) && dp[i][d] >= K - (int)path.size()) {
        path.push_back(x);
        x = e[i][d^1];
        break;
      }
    }
  }
  path.push_back(x);

// sort(path.begin(), path.end());
// fprintf(stderr, "path: ");
//   for(int i:path) fprintf(stderr, "%d, ", i);
// fprintf(stderr, "\n");

  // fprintf(stderr, "free=%d\n", freecount);
  int possible = 0;
  for(int y : path) {
    if (test(y, -1, 0)) ++possible;
  }
  if (!possible) {
    printf("NO\n0\n");
  } else {
    for(int i=0;i<freecount;i++) {
      possible = possible*2%MOD;
    }
    printf("YES\n%d\n", possible);
  }
  
  return 0;
}

Compilation message (stderr)

wells.cpp: In function 'int go(int, int)':
wells.cpp:18:7: warning: unused variable 'x' [-Wunused-variable]
   18 |   int x = e[i][d], y = e[i][d^1];
      |       ^
wells.cpp: In function 'int main()':
wells.cpp:131:21: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  131 |   while(path.size() < K-1) {
      |         ~~~~~~~~~~~~^~~~~
wells.cpp:79:8: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
   79 |   scanf("%d%d", &N, &K);
      |   ~~~~~^~~~~~~~~~~~~~~~
wells.cpp:82:10: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
   82 |     scanf("%d%d", &x, &y);
      |     ~~~~~^~~~~~~~~~~~~~~~

• Subtask 1: 18 / 30
• Subtask 2: 12 / 20
• Subtask 3: 0 / 20
• Subtask 4: 0 / 30