裏紙

ほぼ競プロ、たまに日記

CF 740 D - Alyona and a tree

問題

Problem - D - Codeforces

問題概要

頂点数nの木があり、各頂点には1からnまでの番号が振られている。根の頂点番号は1である。各頂点はそれぞれ整数a_iを持っている。

また、辺には重みがあり、n-1本の辺の情報について、i個目の情報は頂点i+1p_iを結んでおり、その重みはw_iである。

いま、頂点vが頂点u( \neq v)controlしているということを次のように定義する:

  • 頂点uvを根とする部分木に含まれる
  • dist(u,v) \le a_u

各頂点iに対して、その頂点がcontrolしている頂点の個数を求めよ。

  • 1 \le n \le 200000
  • 1 \le a_i \le 10^9
  • 1 \le p_i \le n
  • 1 \le w_i \le 10^9

イデア

controlの条件に関して、uvの距離で定義されているが、今回のグラフは木なので任意の頂点間の距離は根からの距離を利用することで計算できる。根からの各頂点への距離をd_iとすると、controlの定義からuvよりも根から離れていることになるので2番目の条件はd_u - d_v \le a_uとなり、更に変数を揃えるために変形するとd_u - a_u \le d_vということになる。

さて、ここでuを固定して考えてみる。vとしてあり得るのは根からのuへのパス上にある頂点ということになる。根をスタート地点としてd_vは単調増加するので、初めて条件を満たす位置xからuの手前までの区間がこの条件を満たすのでこの区間に+1すればいいということになる。

まず、初めて条件を満たす位置xを探すためには、各頂点からの親をダブリングして保存しておけば二分探索によって探すことが可能になる。

そして、このような区間へのaddはimos法でuの親の頂点に+1,xの親の頂点に-1をしておいて最後に子から親へ足し上げれば答えを出すことが出来る。

実装(C++)

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
#define rep(i,n) for(int (i)=0;(i)<(int)(n);++(i))
#define each(itr,c) for(__typeof(c.begin()) itr=c.begin(); itr!=c.end(); ++itr)
#define all(x) (x).begin(),(x).end()
#define pb push_back
#define fi first
#define se second

typedef pair<int,int> pi;

const int N=200000;
const int LOG_N=18;

struct edge{int to,cost;};

int n;
int a[N];

int par[LOG_N][N];

ll d[N]={};
vector<edge> G[N];
int depth[N]={};
void dfs(int now, int parent)
{
    rep(i,G[now].size())
    {
        edge e = G[now][i];
        if(parent != e.to)
        {
            d[e.to] = d[now]+e.cost;
            depth[e.to] = depth[now]+1;
            dfs(e.to,now);
        }
    }
}

int ans[N]={};

int main()
{
    memset(par,-1,sizeof(par));

    scanf(" %d", &n);
    rep(i,n) scanf(" %d", &a[i]);
    rep(i,n-1)
    {
        int p,w;
        scanf(" %d %d", &p, &w);
        --p;

        G[i+1].pb(edge{p,w});
        G[p].pb(edge{i+1,w});

        par[0][i+1] = p;
    }

    // 根からの距離を計算
    dfs(0,0);

    // ダブリング
    for(int i=1; i<LOG_N; ++i)rep(j,n)
    {
        if(par[i-1][j]!=-1) par[i][j]=par[i-1][par[i-1][j]];
    }

    // 各頂点について、二分探索で条件を満たす境界を探す
    for(int i=1; i<n; ++i)
    {
        int l=0,r=n;
        while(r-l>1)
        {
            int m=(l+r)/2;

            int t=m;
            int now=i;
            for(int j=LOG_N-1; j>=0; --j)
            {
                int sz = 1<<j;
                if(t-sz>=0)
                {
                    if(now!=-1) now = par[j][now];
                    t -= sz;
                }
            }

            bool ok=false;
            if(now!=-1) ok=(d[i]-a[i] <= d[now]);

            if(ok) l=m;
            else r=m;
        }

        int t=l;
        int now=i;
        for(int j=LOG_N-1; j>=0; --j)
        {
            int sz = 1<<j;
            if(t-sz>=0)
            {
                if(now!=-1) now = par[j][now];
                t -= sz;
            }
        }

        if(now!=i)
        {
            ++ans[par[0][i]];
            if(par[0][now]!=-1) --ans[par[0][now]];
        }
    }

    // 子から順に足し上げる
    vector<pi> v(n);
    rep(i,n) v[i]=pi(-depth[i],i);
    sort(all(v));

    rep(i,n-1)
    {
        int now = v[i].se;
        ans[par[0][now]] += ans[now];
    }

    // output
    rep(i,n)
    {
        if(i) printf(" ");
        printf("%d", ans[i]);
    }
    printf("\n");
    return 0;
}