不定期プログラミング覚え書き

青コーダーと黄コーダーの間を彷徨う社会人プロコン勢が余力のあるときに復習した内容をまとめるブログ

yukicoder #363: 門松サイクル

問題

No.363 門松サイクル - yukicoder
すべての連続する3つの数が門松列となっているサイクルを門松サイクルという。
木が与えられた後、2頂点ペアをクエリとして与えるので、もともとの木に2頂点をつなげつことでできるサイクルが
門松サイクルであるか否かを答えよ

考えたこと

一番下にある共通の親を求める時に2^k子上の親をそれぞれ保持しといて見つけるアレ。アレの香りが凄いする(名前を思い出せない)
あとアレと同じ感じで2^kだけ上に辿っても門松列が作れるかどうかを保持してればなんとかなりそう。
でもあまりそれ実装したこと無い気がする……と少し日和る。

いや、そもそもN=10^5のノードが全部縦につながってたら10^5のノードが門松列になってるかの判定するから前処理の時点で10^10使ってLTEだったりする……?なんて一瞬血迷いながら解説をチラ見した。

……やっぱもともとの想定で正しかった。ちなみにアレじゃなくてLCAな。
あと、10^5の一直線ノードが門松列になってるのを判定するのは10^5 * log(10^5)なので問題なかった。

ただし、u,v,LCA(u,v)のサイクルが門松サイクルであると調べるためには、結構色々調べないといけない。
具体的には、(u,v,parent[v]) (parent[u],u,v), LCA(u,v)とその両脇のチェックが必要。あと、2^kの情報を上手くつかいながら u~LCA(u,v)やv~LCA(u,v)が門松列かの判定も必要。
あと、LCA(u,v)がuに等しい時はLCA(u,v)の両脇という概念がないので、特殊処理が必要である。

コード

かなりごちゃごちゃ。滅茶苦茶時間かかったし凄いバグった。
LCAはライブラリ化するにしても、これを本番に解ききる実装力と頭の良さは鍛えていかないとなぁ。

bool kadomatsu(int x,int y, int z){
   if( x==y || y==z || z==x ) return false;
   if( x<y && y>z ) return true;
   if( x>y && y<z) return true;
   return false;
}

int depth[100000];
int parent[18][100000];
int val[100000];
bool iskadsequence[17][100000];
vector<int> edge[100000];

int up(int cur, int d){
   if( d==0 ) return cur;
   int ex=0;
   int ret = cur;
   while( (1<<ex)*2 <= d ) ex++;
   while( d!= 0&&ex>=0 ){
      if( (d&(1<<ex))!=0){
         ret = parent[ex][ret];
         d-=(1<<ex);
      }
      ex--;
   }
   return ret;
}
int lca(int u, int v){
   if( u==v ) return u;
   int left=0;
   int right=depth[u];
   while( left+1 < right ){
      int mid=(left+right)/2;
      if( up(u,mid)==up(v,mid) ){
         right = mid;
      }
      else{
         left=mid;
      }
   }
   return up(u,right);
}

bool checkKadSeq(int u, int d){
   if( d<= 1 ) return true;
   int ex=0;
   while( (1<<ex)*2 <= d ) ex++;
   if( iskadsequence[ex][u]== false ) return false;
   if( d==(1<<ex) ) return true;
   int p = parent[ex][u];
   if( !kadomatsu(val[parent[0][p]],val[p],val[up(u, depth[u]-depth[p]-1)]) ) return false;
   return checkKadSeq(parent[ex][u], d-(1<<ex));
}

int main(){
   int N;
   scanf("%d",&N);
   memset(depth,-1,sizeof(depth));
   memset(parent,-1,sizeof(parent));
   memset( iskadsequence,false,sizeof(iskadsequence));
   for( int i = 0 ; i < N; i++ ) cin >> val[i];
   for( int i = 0 ; i <N-1; i++ ){
      int x,y;
      scanf("%d%d",&x,&y);
      x--;y--;
      edge[x].push_back(y);
      edge[y].push_back(x);
   }
   stack<int> st;
   st.push(0);
   depth[0]=0;
   parent[0][0]=-1;
   while( !st.empty() ){
      int cur = st.top();
      st.pop();
      for( int ei=0; ei < (int)edge[cur].size(); ei++){
         int nxt = edge[cur][ei];
         if( depth[nxt]!=-1  ) continue;
         depth[nxt]=depth[cur]+1;
         parent[0][nxt]=cur;
         iskadsequence[0][nxt]=true;
         int tmpCur=cur;
         int tmpNxt=nxt;
         for( int i = 0;parent[i][tmpCur]!=-1;i++){
            parent[i+1][nxt]=parent[i][tmpCur];
            iskadsequence[i+1][nxt]=(iskadsequence[i][nxt]&iskadsequence[i][tmpCur]&kadomatsu(val[parent[0][tmpCur]],val[tmpCur],val[tmpNxt]));
            tmpCur=parent[i][tmpCur];
            tmpNxt=parent[i][tmpNxt];
         }
         st.push(nxt);
      }
   }
   int Q;
   scanf("%d",&Q);
   for( int i = 0 ; i < Q; i++ ){
      int u,v;
      scanf("%d%d",&u,&v);
      u--;
      v--;
      if( depth[u] > depth[v] ){
         swap(u,v);
      }
      if( parent[0][u]==v || parent[0][v]== u){
         printf("NO\n");
         continue;
      }
      int v2=up(v,depth[v]-depth[u]);
      if( v2 == u ){
         int child_v=up(v,depth[v]-depth[u]-1);
         if( checkKadSeq(v,depth[v]-depth[u])
             && kadomatsu(val[parent[0][v]],val[v],val[u])
             && kadomatsu(val[child_v],val[u],val[v])){
                    printf("YES\n");
         }
         else{
            printf("NO\n");
         }
         continue;
      }
      int p=lca(u,v2);
      int child_u=up(u,depth[u]-depth[p]-1);
      int child_v=up(v,depth[v]-depth[p]-1);
      if( kadomatsu(val[child_u],val[p],val[child_v])
            && checkKadSeq(u,depth[u]-depth[p])
            && checkKadSeq(v,depth[v]-depth[p])
            && kadomatsu(val[u],val[v],val[parent[0][v]])
            && kadomatsu(val[v],val[u],val[parent[0][u]])
            ){
            printf("YES\n");
      }
      else{
            printf("NO\n");
      }
   }
   return 0;        
}