1328E - Tree Queries - CodeForces Solution


dfs and similar graphs trees *1900

Please click on ads to support us..

Python Code:

from collections import defaultdict, deque, Counter
import sys
from decimal import *
from heapq import heapify, heappop, heappush
import math
import random
import string
from copy import deepcopy
from itertools import combinations, permutations, product
from operator import mul, itemgetter
from functools import reduce, lru_cache
from bisect import bisect_left, bisect_right

def input():
    return sys.stdin.readline().rstrip()
def getN():
    return int(input())
def getNM():
    return map(int, input().split())
def getList():
    return list(map(int, input().split()))
def getListGraph():
    return list(map(lambda x:int(x) - 1, input().split()))
def getArray(intn):
    return [int(input()) for i in range(intn)]

mod = 10 ** 9 + 7
MOD = 998244353

inf = float('inf')
eps = 10 ** (-12)
dy = [0, 1, 0, -1]
dx = [1, 0, -1, 0]


class Doubling():
    def __init__(self, n, edges, root):
        self.n = n
        self.logk = n.bit_length()
        self.doubling = [[-1] * self.n for _ in range(self.logk)]
        self.depth = [-1] * N
        self.depth[root] = 0
                par = [-1] * N
        pos = deque([root])
        while len(pos) > 0:
            u = pos.popleft()
            for v in edges[u]:
                if self.depth[v] == -1:
                    self.depth[v] = self.depth[u] + 1
                    par[v] = u
                    pos.append(v)
                for i in range(self.n):
            self.doubling[0][i] = par[i]
        for i in range(1, self.logk):
            for j in range(self.n):
                if self.doubling[i - 1][j] == -1:
                    self.doubling[i][j] = -1
                else:
                    self.doubling[i][j] = self.doubling[i - 1][self.doubling[i - 1][j]]

    def upstream(self, s, t):
        if t == 0:
            return s
        now = s
        for k in range(self.logk):
            if t & (1 << k):
                now = self.doubling[k][now]
        return now

N, M = getNM()
E = [[] for i in range(N)]
for _ in range(N - 1):
    a, b = getNM()
    E[a - 1].append(b - 1)
    E[b - 1].append(a - 1)

D = Doubling(N, E, 0)
for _ in range(M):
    q = getListGraph()
    q = [[D.depth[D.doubling[0][i]], D.doubling[0][i]] for i in q[1:] if i != 0]
    q.sort()
    ans = 'YES'
    for i in range(len(q) - 1, 0, -1):
        d1, p1 = q[i - 1]
        d2, p2 = q[i]
        if D.upstream(p2, d2 - d1) != p1:
            ans = 'NO'
    print(ans)

C++ Code:

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define double long double 
typedef long long LL;
typedef pair<int,int> PII;
template<typename T>
vector<vector<T>> arr(int n,int m){
    return vector<vector<T>>(n,vector<T>(m));
}

const int N=500010,mod=1e9+7;

int n,m;
set<int> g[N];
int d[N],fa[N][32];
int root=1;

void bfs(){
    fill_n(d,N,-1);
    d[0]=0,d[root]=1;
    
    queue<int> q;
    q.push(root);
    
    while(q.size()){
        int t=q.front();
        q.pop();
        
        for(int i:g[t]){
            if(d[i]==-1){
                d[i]=d[t]+1;
                q.push(i);
                
                fa[i][0]=t;
                for(int k=1;k<=31;++k){
                    fa[i][k]=fa[fa[i][k-1]][k-1];
                }
            }
        }
    }
}

int lca(int a,int b){
    if(d[a]<d[b]) swap(a,b);
    for(int i=31;i>=0;--i){
        int x=fa[a][i];
        if(d[x]>=d[b]) a=x;
    }
    if(a==b) return a;
    for(int i=31;i>=0;--i){
        int x=fa[a][i],y=fa[b][i];
        if(x!=y) a=x,b=y;
    }
    return fa[a][0];
}

signed main(){
#ifdef DEBUG
    freopen("test.txt", "r", stdin);
#endif
    ios::sync_with_stdio(0);
    cin.tie(0);
    
    cin>>n>>m;
    for(int i=0;i<n-1;++i){
        int a,b;cin>>a>>b;
        g[a].insert(b);
        g[b].insert(a);
    }
    
    bfs();
    while(m--){
        int k;
        vector<int> t;
        cin>>k;
        for(int i=0;i<k;++i){
            int x;cin>>x;
            t.push_back(x);
        }
        sort(t.begin(),t.end(),[&](int i,int j){
            return d[i]>d[j];
        });
        bool flag=true;
        for(int i=0;i<k-1;++i){
            int a=t[i],b=t[i+1],p=lca(a,b);
            if(!(p==b||g[p].count(b))) flag=false;
        }
        cout<<(flag?"YES":"NO")<<'\n';
    }
}


Comments

Submit
0 Comments
More Questions

1615B - And It's Non-Zero
1619E - MEX and Increments
34B - Sale
1436A - Reorder
1363C - Game On Leaves
1373C - Pluses and Minuses
1173B - Nauuo and Chess
318B - Strings of Power
1625A - Ancient Civilization
864A - Fair Game
1663B - Mike's Sequence
448A - Rewards
1622A - Construct a Rectangle
1620A - Equal or Not Equal
1517A - Sum of 2050
620A - Professor GukiZ's Robot
1342A - Road To Zero
1520A - Do Not Be Distracted
352A - Jeff and Digits
1327A - Sum of Odd Integers
1276A - As Simple as One and Two
812C - Sagheer and Nubian Market
272A - Dima and Friends
1352C - K-th Not Divisible by n
545C - Woodcutters
1528B - Kavi on Pairing Duty
339B - Xenia and Ringroad
189A - Cut Ribbon
1182A - Filling Shapes
82A - Double Cola