CENS20F - Editorial

Problem Link:

Practice
Contest

Author & Editorialis: Arun Das
Tester: Jatin Nagpal

DIFFICULTY:

EASY-MEDIUM

PREREQUISITES:

DFS(Depth First Search)

PROBLEM:

Given a rooted tree and Q queries, in each query transfer the values of nodes at a non-zero even distance to the source node. Print the final values of the nodes.

QUICK EXPLANATION:

Create a N\times2 2D array to keep track of whether nodes at all even/odd distances from a node have been visited. Start a DFS from the source node in each query and if all even/odd nodes have already been visited skip the entire subtree.

EXPLANATION:

Notice that once we complete the DFS of a source node, all values at non-zero even distances in its subtree will become zero and will not change further.

Let’s see how the search goes for each query.

  • First, we start looking for nodes at an even distance from the source. (The source value is not added to itself)
  • As the DFS continues, we look for odd distanced nodes from the child of the source.
  • Next, we look for even distanced nodes from the grandchildren of the source. (These values are added to the source)

We see that all nodes in a subtree can be divided into 2 parts - Nodes that are at an even distance from the root of the subtree and nodes that are at an odd distance. So let’s create an N\times2 size 2 dimensional visited array to keep track of this. visited[i][0] will be zero if all nodes in the subtree of i that are at non-zero even distances from i are 0 and visited[i][1] will keep track of the same for odd distanced nodes.

Now when we are doing the DFS and looking for even distanced nodes from the current node, we check visited[current][0]. If this value is 0, it means that all nodes at even distances haven’t been visited and we continue the DFS. Otherwise, all nodes at an even distance from this node have been visited and we can just add the value of the current node to the source (if it’s not the source itself) and skip the entire subtree.

Why add if visited?

Consider the following tree and the queries -> 5, 4, 2 in order.
graph

After the first and second queries, visited[5][0] and visited[4][0] will become 1 but note that the values of the source itself may not become 0. When we query for 2, we need to add the values of nodes 4 and 5.

The time complexity of this solution in O(n).

Why?

Each node will be accessed at most 3 times.

  • Once when this node is the query. After the query visited[node][0] will become 1.
  • Once when a query starts from an ancestor of this node which is at an even distance. After this query, the value of this node will become 0.
  • Once when a query starts from an odd distanced ancestor of this node. After this query visited[node][1] will become 1.
    Any further queries will either terminate in one of its ancestors or can be skipped entirely.
Things to watch out for
  • Make sure to mark the source node visited or you might TLE in case the input is star-shaped.
  • Use the appropriate data type since the answer might be as big as about 1014

This can be implemented in more than one ways and there can be other different approaches. Feel free to discuss them. :slight_smile:

SOLUTIONS:

Setter's Solution(with comments)
#include<bits/stdc++.h>
using namespace std;

#define int long long

int visited[200005][2];
int values[200005], parent[200050];

void dfs(int vertex, int parity, vector<int> tree[], int source, int parent){
    //if the vertex is not source and distance from source is even add the value
    if(vertex != source && parity == 0){
        values[source] += values[vertex];
        values[vertex] = 0;
    }

    //if visited[vertex][1] == 1, all the nodes at odd distance
    //from the current vertex are zero. So there is no  need to go further.
    if(visited[vertex][parity] == 0){
        for(auto x: tree[vertex]){
            if(x != parent){
                //if we are looking for nodes at an even distance from current node,
                //then we have to look for nodes at an odd distance from this node's children
                dfs(x, parity^1, tree, source, vertex);
            }
        }

        //if parity == 1, all nodes at odd distance from this vertex have been visited
        //if parity == 0, all nodes at even distance from this vertex have been visited
        visited[vertex][parity] = 1;
    }

}

void findParents(int v, vector<int> tree[], int par){
    parent[v] = par;
    for(auto x: tree[v])
        if(x != par)
            findParents(x, tree, v);
}

int32_t main(){
    // ios_base::sync_with_stdio(false);
    // cin.tie(NULL);
    
    int t;
    cin >> t;
    while(t--){
        int n, i, j, x, u, v, q;
        cin >> n >> q;

        for(i=0; i<200005; i++){
            visited[i][0] = 0;
            visited[i][1] = 0;
        }

        // take the input values
        for(i=0; i<n; i++)
            cin >> values[i];

        //take the input tree
        vector<int> tree[n];
        for(i=0; i<n-1; i++){
            cin >> u >> v;
            u--; v--;
            tree[u].push_back(v);
            tree[v].push_back(u);
        }

        // find the parent of each vertex
        findParents(0, tree, -1);
        while(q--){
            cin >> x; x--;
            if(visited[x][0] == 1){
                continue;
            }
            dfs(x, 0, tree, x, parent[x]);
            visited[x][0] = 1;

        }

        for(i=0; i<n; i++){
            cout << values[i] << " ";
        }
        cout << "\n";
    }
} 
Tester's Solution
#include <bits/stdc++.h>
using namespace std;
#define ff first
#define ss second
#define MP make_pair
#define PB push_back
#define ll long long
#define int long long
#define f(i,x,n) for(int i=x;i<n;i++)
#define ld long double
const int mod=1000000007;
const int INF=1e18;
int n,q,a[200005],vis[200005],ans[200005];
vector <int> gp[200005];
pair <int,int> dfs(int i,int par,int state) {
    int flag=0;
    if(vis[i]==1 && (state&1) ==0)
    {
        state|=1;
        flag=1;
    }
    pair <int,int> ret={0,0};
    for(auto j: gp[i]) {
        if(j!=par) {
            pair <int,int> tmp=dfs(j,i, (state%2)*2 + state/2 );
            ret.ff+=tmp.ff;
            ret.ss+=tmp.ss;
        }
    }
    ret.ff+=a[i];
    if(flag==1)
        ans[i]=ret.ff;
    else
        ans[i]=(1^(state&1))*a[i];
    ret.ff=ret.ff^ret.ss;
    ret.ss=ret.ff^ret.ss;
    ret.ff=ret.ff^ret.ss;
    return ret;
}
int32_t main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    int t;
    cin>>t;
    while(t--) {
        cin>>n>>q;
        f(i,1,n+1)
        {
            vis[i]=0;
            gp[i].clear();
            ans[i]=0;
            cin>>a[i];
        }
        f(i,0,n-1) {
            int u,v;
            cin>>u>>v;
            gp[u].push_back(v);
            gp[v].push_back(u);
        }
        while(q--) {
            int Q;
            cin>>Q;
            vis[Q]=1;
        }
        dfs(1,0,0);
        f(i,1,n+1)
            cout<<ans[i]<<" ";
        cout<<'\n';
    }
    
    return 0;
}
7 Likes

Thanks for providing such a neat code, setters rarely provide such a readable and documented code, btw idea of squeezing everything in a single dfs is really cool

MY solution
#include <bits/stdc++.h>
using namespace std;
#define int long long 
#define pb push_back 
#define vi vector
#define FOR(i,l,n) for(int i=l;i<n;i++)
#define EACH(x,a) for(auto x:a)
const long mxN =2e5+2 ;
int n,q,a[mxN],a1[mxN],a2[mxN] ,d[mxN],ans[mxN];
vi<int> adj[mxN];
void dfs1(int v=1,int p=0){
  if(v>1)d[v]=d[p]+1 ;
  a2[v]=d[v]&1?a[v]:0 ;a1[v]=d[v]&1?0:a[v] ;
  EACH(x,adj[v])
    if(x^p)
      dfs1(x,v),a1[v]+=a1[x],a2[v]+=a2[x];
}
set<int>e,o ;
void dfs2(int v=1,int p=0,bool ok1=0,bool ok2=0){
  if(d[v]&1^1){
    if(ok1)
      ans[v]=0;
    else{
      ans[v]=a[v] ;
      if(e.count(v))
        ok1=1 ,ans[v]=a1[v] ;
    }    
  }else{
    if(ok2)
      ans[v]=0;
    else{
      ans[v]=a[v] ;
      if(o.count(v))
        ok2=1 ,ans[v]=a2[v] ;
    }    
  }
  EACH(x,adj[v])
    if(x^p)
      dfs2(x,v,ok1,ok2) ;
}
void solve(){ 
  cin >>n>>q ;
  FOR(i,1,n+1)
    cin >> a[i] ;
  FOR(i,0,n-1){
    int a,b ;cin>>a>>b ;
    adj[a].pb(b);adj[b].pb(a) ;
  }
  dfs1() ;
  while(q--){
    int x;cin >> x ;
    if(d[x]&1)
      o.insert(x);
    else
      e.insert(x);
  }
  dfs2() ;
  FOR(i,1,n+1)
    cout << ans[i] <<" " ;
  cout << "\n" ;
  FOR(i,1,n+1)
    adj[i].clear(),d[i]=0,ans[i]=0 ;
  e.clear();o.clear() ;
}
signed main() {
  int t;cin >>t ;
  while(t--)
    solve();
}
1 Like

What is wrong with this ->https://www.codechef.com/viewsolution/37003321
Please help someone .I used BFS and used checked array to prevent from visiting tree with all 0s.

I wrote a solution but got stuck since I couldn’t see what went wrong here. Can someone provide a simple test case so that I realize my mistake here?

My approach: Calculate the depth of each node using a BFS (considering root to be at depth 0), and then create a list of each node corresponding to it’s depth, so that the list nodeAtDepth[i] tells us what nodes are there at depth i.

Then, for each query V, consider all nodes k at distance depth[V] + 2y and add A[k] to A[V], and set A[k] to 0. If I find a depth which I’ve already visited before, all of those values must have been set to 0 and so I can discard it.

Link to my solution

We can created a separate forest from the given tree, with edges connecting node to its grand parent.

Can you share the editorial of Cherry and Squares

It is not fully prepared. Will post it by tomorrow. :slight_smile:

Consider this case
1
8 1
1 1 1 1 1 1 1 1
1 2
2 3
3 4
4 7
2 5
5 6
6 8
5
If two nodes are at different even depths it doesn’t necessarily mean that one is in the subtree of the other.
Your output - 1 1 1 1 3 1 0 0
Correct output - 1 1 1 1 2 1 1 0

3 Likes

Okay thanks.

Don’t know why getting Tle … i think its a O(n) solution. ti would be great help if someone help me finding the reason .
https://www.codechef.com/viewsolution/37002710

Oh dang, that was stupid of me.
Thank you for taking the time to go through it!
I’ll upsolve this tomorrow.

Hello, can you tell me on which test-case does my solution fails ?
https://www.codechef.com/viewsolution/37006323
Same idea as that of the editorial.
@aert

@aert @cub Can anyone please help me find where my solution fails :https://www.codechef.com/viewsolution/37006775
@aryanc403

1 Like

I haven’t properly seen your code but it fails on
1
8 3
1 1 1 1 1 1 1 1
1 2
2 3
3 4
4 7
2 5
5 6
6 8
7
3
1
Your output - 4 1 0 1 0 1 0 1
Correct output - 5 1 0 1 0 1 0 0

I was coding segment tree along with euler tour. Only Update function was left when I realised one doesn’t need it. :slightly_smiling_face:
“Things to watch out for” Did got one TLE due to this.

2 Likes

What is wrong with this -> https://www.codechef.com/viewsolution/37006904
Please help someone .I used BFS and created a distance array for each Node.

I am getting runtime array but it is working fine for the sample test case.

Thanks brother , I was missing a Q.pop() operation when I was using continue statement But after correcting it also I am getting correct answer on your test case but WA during submission. Can you now tell me some test cases where my solution is failing . Thanks in advance.
https://www.codechef.com/viewsolution/37007493

Actually your approach is not correct. We have to transfer the values of nodes at even distance to the source node “present in that subtree”.
But according to your approach you are transfering values of all nodes whose depth is even from that node, so there may be some nodes which can not be present in that subtree.

1 Like

" You are given a tree rooted at node 1 with N vertices."
This is the first line of the question .

3 Likes

Yeah, @aert helped me realize it as well.
Thank you!