Passeio

Passeio.cpp

// Passeio - Seletiva IOI 2013
// Complexidade O(N*log²(N))
// Matheus Leal V

#include <bits/stdc++.h>
#define mid ((a + b)/2)
#define inf 10000000
#define N 100050
#define f first
#define s second
using namespace std;
typedef pair<int, int> pii;
typedef long long ll;

pii best;

int n, A, B, sz[N], tot, qtd[N], block[N];

vector<int> grafo[N];

ll resp;

struct node
{
    int val;
    node *l, *r;
    node()
    {
        val = 0;
        l = r = NULL;
    }
};

node *root;

int Val(node *root)
{
    return root != NULL ? root->val : 0;
}

void upd(node *root, int a, int b, int i)
{
    if(a == b)
    {
        root->val ++;
        
        return;
    }
    if(i <= mid)
    {
        if(root->l == NULL) root->l = new node();

        upd(root->l, a, mid, i);
    }
    else
    {
        if(root->r == NULL) root->r = new node();

        upd(root->r, mid + 1, b, i);
    }

    root->val = Val(root->l) + Val(root->r);
}

int query(node *root, int a, int b, int i, int j)
{
    if(j < a || i > b) return 0;

    if(i <= a && j >= b) return root->val;

    int A = (root->l != NULL ? query(root->l, a, mid, i, j) : 0);

    int B = (root->r != NULL ? query(root->r, mid + 1, b, i, j) : 0);

    return A + B;
}

void preencher(int x, int p, int dist)
{
    upd(root, 0, N, dist);

    for(int i = 0; i < grafo[x].size(); i++)
    {
        int v = grafo[x][i];

        if(v == p || block[v]) continue;

        preencher(v, x, dist + 1);
    }
}

void getans(int x, int p, int dist)
{
    if(dist > B) return;

    if(A <= dist && dist <= B) resp ++;

    resp += (ll) query(root, 0, N, max(A - dist, 0), max(B - dist, 0));

    for(int i = 0; i < grafo[x].size(); i++)
    {
        int v = grafo[x][i];

        if(block[v] || v == p) continue;

        getans(v, x, dist + 1);
    }
}

int tam(int x, int p)
{
    sz[x] = 1;

    for(int i = 0; i < grafo[x].size(); i++)
    {
        int v = grafo[x][i];

        if(v == p || block[v]) continue;

        sz[x] += tam(v, x);
    }

    return sz[x];
}

void split(int x, int p)
{
    int maior = tot - sz[x];

    for(int i = 0; i < grafo[x].size(); i++)
    {
         int v = grafo[x][i];

         if(v == p || block[v]) continue;

         split(v, x);

         maior = max(maior, sz[v]);
    }

    if(maior < best.f) best = pii(maior, x);
}

int find_centroid(int x)
{
    tot = tam(x, x);

    best = pii(inf, 0);

    split(x, x);

    return best.s;
}

int decomp(int x, int nivel)
{
     x = find_centroid(x);

    block[x] = 1;

    root = new node();

    for(int i = 0; i < grafo[x].size(); i++)
    {
        int v = grafo[x][i];

        if(!block[v])
        {
            getans(v, x, 1);

            preencher(v, x, 1);
        }
    }

    for(int i = 0; i < grafo[x].size(); i++)
    {
        int v = grafo[x][i];

        if(block[v]) continue;

        decomp(v, nivel + 1);
    }

    return x;
}

 main()
{
    ios::sync_with_stdio(false); cin.tie(0);

    cin>>n;

    for(int i = 1, a, b, c; i < n; i++)
    {
        cin>>a>>b;

        a++, b++;

        grafo[a].push_back(b);

        grafo[b].push_back(a);
    }

    cin>>A>>B;

    decomp(1, 1);

    cout<<resp<<"\n";
}