Submission #409481

#	Time	Username	Problem	Language	Result	Execution time	Memory
409481		pauloamed	Railway (BOI17_railway)	C++14	100 / 100	750 ms	135020 KiB

railway

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

#define MAXN 500010
#define MAXLOG 21

// global time var for computing the following arrays
// (only incremented when entering node)
int currTime;
// inTime[v]: time in which node v was entered
// outTime[v]: time in which node v was left
// time2node[t]: node which was entered in time t
// inTime[time2Node[t]] = t // time2Node[inTime[v]] = v
int inTime[MAXN], outTime[MAXN], time2Node[MAXN];
int sz[MAXN]; // sz[v]: size of subtree rooted at v
vector<int> v[MAXN]; // adj list
int depth[MAXN];

vector<int> empt[MAXN];
void precalc(int x, int par, int d){
  time2Node[currTime] = x; inTime[x] = currTime++;
  depth[x] = d;
  sz[x] = empt[x].size() + 1;
  for(auto y : v[x]){
    if(y == par) continue;
    precalc(y, x, d + 1);
    sz[x] += sz[y];
  }
  outTime[x] = currTime;
}

bool isInSubtreeOf(int x, int y){
  return inTime[x] >= inTime[y] && outTime[x] <= outTime[y];
}

// pra cada conjunto de nos eu guardo que ali comeca a query q
vector<int> queriesStart[MAXN];

// pra cada LCA de um conjunto de nos eu guardo que ali acaba a query i
vector<int> queriesEnd[MAXN];

int ans[MAXN];
int timeAns[MAXN];

map<int,int> on;
set<int> off;

map<pair<int,int>, int> edge2id;

void solve(int x, int par, bool keep){
  // retrieving heavy child from x
  int heavyChildSize = -1;
  int heavyChild = -1;
  for(auto y : v[x]){
    if(y == par) continue;
    if(sz[y] > heavyChildSize){
      heavyChild = y;
      heavyChildSize = sz[y];
    }
  }

  // vising light and then heavy childs of x
  for(auto y : v[x]){
    if(y == heavyChild || y == par) continue;
    solve(y, x, false);
  }
  if(heavyChild != -1) solve(heavyChild, x, true);

  for(auto y : v[x]){
    if(y == heavyChild || y == par) continue;
    for(int t = inTime[y]; t < outTime[y]; ++t){
      int z = time2Node[t];
      for(auto y : queriesStart[z]) on[y]++;
      for(auto y : queriesEnd[z]) off.insert(y);
    }
  }

  for(auto y : queriesStart[x]) on[y]++;
  for(auto y : queriesEnd[x]) off.insert(y);

  // cout << "========================\n";
  // cout << x << "\n";
  // for(auto y : on) cout << y.first << " " << y.second << ", "; cout << "\n";
  // for(auto y : off) cout << y << " "; cout << "\n";
  // cout << "========================\n";

  ans[edge2id[{x, par}]] = on.size() - off.size();

  if(!keep){
    for(int t = inTime[x]; t < outTime[x]; ++t){
      int z = time2Node[t];
      for(auto y : queriesEnd[z]){
        off.erase(y);
      }
      for(auto y : queriesStart[z]){
        on[y]--;
        if(on[y] == 0) on.erase(y);
      }
    }
  }
}


int st[MAXN][MAXLOG]; // sparse table (x,i) guardando ancestral de ordem 2^i de x
bool vis_lca[MAXN]; // lvl=profundidade de cada no

void init_dfs(int x, int par){
    vis_lca[x] = true; // marcando vertice como visitado
    st[x][0] = par; // o ancestral de ordem 2^0 (1) de x eh o pai dele
    // se tiver pai valido (!= raiz), a prof aumenta em 1
    if(par != -1) depth[x] = depth[par] + 1;

    for(int i = 0; i < v[x].size(); ++i) // visitar todos adjacentes
        if(!vis_lca[v[x][i]]) init_dfs(v[x][i], x); // passo recursivo
}

void init_st(){
    // processo cada nivel de cada vertice
    for(int x = 1; x < MAXLOG; ++x) // 2^0 ja foi calculado, calculo pro resto
        for(int i = 0; i < MAXN; ++i){ // pra cada vertice
            // olho o ancestral 2^(i-1) do meu 2^(i-1), achando entao o meu 2^i
            if(st[i][x-1] == -1) st[i][x] = -1;
            else st[i][x] = st[st[i][x-1]][x-1];
    }
}

void init(int root, int n){ // funcao init, raiz e #nos
    // botando que o pai da raiz eh -1, pode ser ela mesma tbm
    init_dfs(root, -1); // calcular os ancestrais imediatos
    init_st(); // init da sparse table
}


int lca(int x, int y){ // lca de x e y
    // cout << lvl[x] << " " << lvl[y] << endl;
    if(depth[x] < depth[y]) swap(x, y); // quero q x seja mais profundo q y

    int falta_subir = (depth[x] - depth[y]); // igualo as profundidades
    // simples representacao binaria de (falta_subir)
    for(int i = MAXLOG-1; i >= 0; --i){ // encontro os bits para representar o
        if((1<<i) <= falta_subir){ //      numero, dos mais signif. para os menos
            falta_subir -= (1<<i);
            x = st[x][i];
        }
    }

    if(x == y) return x; // ocorre quando x ta numa subarvore de y

    // acho o ponto abaixo do encontro (LCA)
    // se eu tentar subir 2^i e eles ja estiverem juntos, nao subo
    // tento entao subir 2^(i-1)
    for(int i = MAXLOG-1; i >= 0; --i){
        if(st[x][i] != st[y][i]){ // se continuarem diferentes, subo
            x = st[x][i]; // subindo pra x
            y = st[y][i]; // subindo pra y
        }
    }

    return st[x][0]; // retornando o ponto de encontro
}

int main(){
  ios::sync_with_stdio(false);
  int n, m, k; cin >> n >> m >> k;

  for(int i = 1; i < n; ++i){
    int a, b; cin >> a >> b; a--, b--;
    v[a].push_back(b);
    v[b].push_back(a);
    edge2id[{a, b}] = i;
    edge2id[{b, a}] = i;
  }

  init(0, n);
  precalc(0, -1, 0);

  for(int i = 0; i < m; ++i){
    int s; cin >> s;
    vector<int> nodes(s);
    for(auto &x : nodes){
      cin >> x; x--;
    }
    int lcax = nodes[0];
    for(int i = 1; i < s; ++i){
      lcax = lca(lcax, nodes[i]);
    }
    // cerr << i << " " << lcax << endl;
    for(auto x : nodes){
      if(x == lcax) continue;
      queriesStart[x].push_back(i);
    }
    queriesEnd[lcax].push_back(i);
  }


  solve(0, -1, true);

  vector<int> anss;
  for(int i = 1; i <= n; ++i){
    // cout << i << " " << ans[i] << endl;
    if(ans[i] >= k) anss.push_back(i);
  }

  cout << anss.size() << "\n";
  for(auto x : anss) cout << x << " ";
  cout << "\n";
}

Compilation message (stderr)

railway.cpp: In function 'void init_dfs(int, int)':
railway.cpp:113:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  113 |     for(int i = 0; i < v[x].size(); ++i) // visitar todos adjacentes
      |                    ~~^~~~~~~~~~~~~

• Subtask 1: 8 / 8
• Subtask 2: 15 / 15
• Subtask 3: 7 / 7
• Subtask 4: 29 / 29
• Subtask 5: 16 / 16
• Subtask 6: 25 / 25