TREEDISTSET - Editorial

editorial
, ,
1

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: sushil2006
Tester: jay_1048576
Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

Diameter of a tree

PROBLEM:

You’re given a tree on N vertices.
Color each of its vertices either red or blue such that:

  • For any integer d, if there exists a pair of vertices in the tree with distance d, there should exist such a pair with the vertices having different colors.

Among all such colorings, minimize the difference between the number of red and blue vertices.

EXPLANATION:

First, we’ll try to satisfy the condition on distances, i.e, achieve every distance using two vertices of opposite colors.

Let D denote the diameter of the tree, i.e, the longest path in it. Let x_D and y_D be two endpoints of a diameter.
The distances we need to worry about are then 1, 2, \ldots, D.
Finding D (and the vertices x_D and y_D) is a classical problem, and can be done in linear time.

How?

This blog is a good read if you’re new to the concept of a tree’s diameter.

For now, let’s focus on just the path from x_D to y_D (which is a case we need to solve anyway, for instance when the tree is a line).
Note that this path contains every distance from 1 to D already - in fact, a very simple construction would be to color x_D red and all the other vertices blue.

However, recall that we want to minimize the difference between the number of red and blue vertices used; so using only one red vertex doesn’t seem optimal: ideally, we’d have about half each of red and blue.

So, we simply generalize our construction a bit: instead of one red vertex and everything else blue, simply color the first half of them red, and the other half blue!
It’s not hard to see that the distance condition is still true.

Proof

For convenience, let’s label the vertices on the path 1, 2, \ldots, D+1 in order.

Suppose vertices 1, 2, 3, \ldots, k are colored red and the rest are blue; for some 1\leq k\leq D.
Then,

  • The distances from vertex 1 to the blue vertices are k, k+1, k+2, \ldots, D.
  • The distances from vertex k+1 to the red vertices are 1, 2, 3, \ldots, k.
  • Together, they cover everything from 1 to D, as required.

This means we’re free to choose any k we like; so of course we can choose k = \left\lfloor\frac{D}{2}\right\rfloor

Now, just by coloring the diameter, we’ve ensured that:

  • All the distances from 1 to D are taken care of.
  • The difference between the number of red and blue vertices is either 0 or 1 (depending on the parity of D), which is the best we can do.

This means all the non-diameter vertices can be freely colored without having to worry about the distance condition at all; so we can use them to ensure the red-blue difference never gets large.
That is, for each uncolored vertex, do the following:

  • If |R-B| = 1, color it red or blue appropriately to bring the difference to 0.
  • If |R-B| = 0, give it any color.

TIME COMPLEXITY:

\mathcal{O}(N) per testcase.

CODE:

Author's code (C++) 
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

template<typename T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long long int ll;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;

#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL)
#define pb push_back
#define endl '\n'
#define sz(a) (int)a.size()
#define setbits(x) __builtin_popcountll(x)
#define ff first
#define ss second
#define conts continue
#define ceil2(x,y) ((x+y-1)/(y))
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define yes cout << "Yes" << endl
#define no cout << "No" << endl

#define rep(i,n) for(int i = 0; i < n; ++i)
#define rep1(i,n) for(int i = 1; i <= n; ++i)
#define rev(i,s,e) for(int i = s; i >= e; --i)
#define trav(i,a) for(auto &i : a)

template<typename T>
void amin(T &a, T b) {
    a = min(a,b);
}

template<typename T>
void amax(T &a, T b) {
    a = max(a,b);
}

#ifdef LOCAL
#include "debug.h"
#else
#define debug(x) 42
#endif

/*



*/

const int MOD = 1e9 + 7;
const int N = 2e3 + 5;
const int inf1 = int(1e9) + 5;
const ll inf2 = ll(1e18) + 5;

vector<ll> adj[N];
pll best;

void dfs1(ll u, ll p, ll d){
    pll px = {d,u};
    amax(best,px);
 
    trav(v,adj[u]){
        if(v == p) conts;
        dfs1(v,u,d+1);
    }
}
 
vector<ll> curr_path;
vector<ll> path;
 
void dfs2(ll u, ll p, ll des){
    curr_path.pb(u);
    if(u == des){
        path = curr_path;
    }
    trav(v,adj[u]){
        if(v == p) conts;
        dfs2(v,u,des);
    }
    curr_path.pop_back();
}

void solve(int test_case)
{
    ll n; cin >> n;
    rep1(i,n){
        adj[i].clear();
    }

    rep1(i,n-1){
        ll u,v; cin >> u >> v;
        adj[u].pb(v), adj[v].pb(u);
    }
 
    best = {-1,-1};
    dfs1(1,-1,0);
    ll s = best.ss;
    best = {-1,-1};
    dfs1(s,-1,0);
    auto [diam,t] = best;
    diam++;
    
    dfs2(s,-1,t);
    vector<ll> col(n+5);
    ll r = 0, b = 0;

    rep(i,sz(path)){
        ll u = path[i];
        if(i < diam/2){
            col[u] = 1;
            r++;
        }
        else{
            col[u] = 2;
            b++;
        }
    }

    vector<bool> on_path(n+5);
    trav(u,path) on_path[u] = 1;

    rep1(u,n){
        if(on_path[u]) conts;
        if(r <= b){
            col[u] = 1;
            r++;
        }
        else{
            col[u] = 2;
            b++;
        }
    }

    assert(abs(r-b) <= 1);

    rep1(i,n){
        if(col[i] == 1) cout << 'R';
        else cout << 'B';
    }

    cout << endl;
}

int main()
{
    fastio;

    int t = 1;
    cin >> t;
    
    rep1(i, t) {
        solve(i);
    
    }
    return 0;
}
Tester's code (C++) 
/*...................................................................*
 *............___..................___.....____...______......___....*
 *.../|....../...\........./|...../...\...|.............|..../...\...*
 *../.|...../.....\......./.|....|.....|..|.............|.../........*
 *....|....|.......|...../..|....|.....|..|............/...|.........*
 *....|....|.......|..../...|.....\___/...|___......../....|..___....*
 *....|....|.......|.../....|...../...\.......\....../.....|./...\...*
 *....|....|.......|../_____|__..|.....|.......|..../......|/.....\..*
 *....|.....\...../.........|....|.....|.......|.../........\...../..*
 *..__|__....\___/..........|.....\___/...\___/.../..........\___/...*
 *...................................................................*
 */
 
#include <bits/stdc++.h>
using namespace std;
#define int long long

void dfs(vector<int> adj[],int dis[],int par[],int u,int p)
{
    par[u]=p;
    for(auto v:adj[u])
    {
        if(v!=p)
        {
            dis[v]=dis[u]+1;
            dfs(adj,dis,par,v,u);
        }
    }
}

void solve(int tc)
{
    int n;
    cin >> n;
    vector<int> adj[n];
    for(int i=0;i<n-1;i++)
    {
        int u,v;
        cin >> u >> v;
        u--,v--;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    int dis[n],par[n];
    dis[0]=0;
    dfs(adj,dis,par,0,0);
    int r=0;
    for(int i=0;i<n;i++)
        if(dis[r]<dis[i])
            r=i;
    dis[r]=0;
    dfs(adj,dis,par,r,r);
    int l=r;
    for(int i=0;i<n;i++)
        if(dis[l]<dis[i])
            l=i;
    vector<int> diameter;
    int u=l;
    while(u!=r)
    {
        diameter.push_back(u);
        u=par[u];
    }
    diameter.push_back(r);
    string ans;
    for(int i=0;i<n;i++)
        ans+='?';
    int red=0,blue=0;
    for(int i=0;i<(diameter.size())/2;i++)
    {
        ans[diameter[i]]='R';
        red++;
    }
    for(int i=(diameter.size())/2;i<diameter.size();i++)
    {
        ans[diameter[i]]='B';
        blue++;
    }
    for(int i=0;i<n;i++)
    {
        if(ans[i]=='?')
        {
            if(red<blue)
            {
                ans[i]='R';
                red++;
            }
            else
            {
                ans[i]='B';
                blue++;
            }
        }
    }
    cout << ans << '\n';
}

int32_t main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
    int tc=1;
    cin >> tc;
    for(int ttc=1;ttc<=tc;ttc++)
        solve(ttc);
    return 0;
}
Editorialist's code (Python) 
def bfs(g, src):
    n = len(g)
    dist = [n]*n
    dist[src] = 0
    qu = [src]
    for u in qu:
        for v in g[u]:
            if dist[v] == n:
                dist[v] = 1 + dist[u]
                qu.append(v)
    return dist

for _ in range(int(input())):
    n = int(input())
    g = [ [] for _ in range(n) ]
    for _ in range(n-1):
        u, v = map(int, input().split())
        g[u-1].append(v-1)
        g[v-1].append(u-1)
    
    dist0 = bfs(g, 0)
    dx = max(list(range(n)), key=dist0.__getitem__)
    distx = bfs(g, dx)
    dy = max(list(range(n)), key=distx.__getitem__)
    disty = bfs(g, dy)
    diameter = distx[dy]
    
    col = ['A']*n
    for i in range(n):
        if distx[i] + disty[i] == diameter:
            if 2*distx[i] <= diameter: col[i] = 'R'
            else: col[i] = 'B'
    dif = col.count('R') - col.count('B')
    for i in range(n):
        if col[i] != 'A': continue
        if dif > 0:
            col[i] = 'B'
            dif -= 1
        else:
            col[i] = 'R'
            dif += 1
    print(''.join(col))
1 Like
2

Why the constraints allows quadratics solutions?

2 Likes
3

I guess the constraints are low (2000) so that the solution can be judged whether our assignment of colors are correct or not with n^2 algorithm checker.(by pairing each node with other and finding it’s distance).

1 Like
4

this actually kind of makes sense, thanks for the insight.

7

Can Anyone tell me why it is giving run time err-#include<bits/stdc++.h>
include
using namespace std;

using ll = long long;

define yes std::cout << “YES” << std::endl
define no std::cout << “NO” << std::endl

define for_n(n) for (int i = 0; i < n; ++i)

define v(n)
std::vector v(n);
for (int i = 0; i < n; ++i) {
std::cin >> v[i];
}

int x=0,maxd=0;
vectormaxpath;
void dfs(int node,int par,vectoradj,int cnt,vector& path){
if(cnt > maxd){
maxd=cnt;
maxpath=path;
x=node;
}

for(auto it:adj[node]){
    if(it != par){
        path.push_back(it);
        dfs(it,node,adj,cnt+1,path);
        path.pop_back();
    }
}

}

void solve(){
int n;
cin >> n;
vectoradj[n+1];
for(int i=0;i<n-1;i++){
int a,b;
cin>>a>>b;
adj[a].push_back(b);
adj[b].push_back(a);
}
// diameter work
vectorpath;
path.push_back(1);
dfs(1,0,adj,0,path);
int y=x;
x=0,maxd=0;
path.clear();
maxpath.clear();
path.push_back(y);
dfs(y,0,adj,0,path);

//  diameter nodes are x and y with maxd and path stored in maxpath
// cout<<x<<" "<<y<<" "<<maxd<<endl;

// finding total diameter nodes
int total=maxpath.size();
int rem=n-total;

// cout<<total<<" "<<rem<<endl;

// choose red and blue from dia nodes
int red,blue;
if(total%2) red=(total+1)/2;
else red=total/2;
blue=total/2;

// choose red and blue from remaining nodes
    blue+=(rem+1)/2;
    red+=rem/2;
    
    // cout<<red<<" "<<blue<<endl;
    for(int i=0;i<red;i++) cout<<'R';
    for(int i=0;i<blue;i++) cout<<'B';
    cout<<endl;

}

int main() {
int t;
cin >> t;
while(t–){
solve();
}

return 0;

}

8

Similar problem from a recent CF contest: Problem - C - Codeforces