2019 NanChan ICPC K.Tree

求

• $v[i]+v[j]=2v[lca(i,j)]$
• $dis(i,j)\leq k$
• $i,j$不在一条链上。

由于条件$3$，无法用点分治。

考虑$dsu \ on \ tree$。

只需要每次计算单独链计算，并且计算完添加进去。

开个$Segment$，$rt[i]$表示$v[x]=i$的所有$dep$，动态开点，动态删除，加个回收就$ok$了。

一遍过。

代码

#include <bits/stdc++.h>
#include <set>
using namespace std;
typedef long long ll;
#define pii pair<int, int>
#define mk make_pair
const int N = 1e6 + 10;
const int mod = 1e9 + 7;

int read()
{
    int x = 0, f = 1;
    char c = getchar();
    while (c < '0' || c > '9')
    {
        if (c == '-')
            f = -1;
        c = getchar();
    }
    while (c >= '0' && c <= '9')
        x = (x << 1) + (x << 3) + c - '0', c = getchar();
    return x * f;
}
ll res = 0;
struct Segment
{
    int ls[N << 4], rs[N << 4], siz[N << 4];
    queue<int> q;
    int scnt = 0;
    int Newpos()
    {
        int x;
        if (!q.empty())
        {
            x = q.front();
            q.pop();
        }
        else
            x = ++scnt;
        return x;
    }
    void pushup(int pos)
    {
        siz[pos] = siz[ls[pos]] + siz[rs[pos]];
    }
    void insert(int &pos, int w, int l, int r)
    {
        if (!pos)
            pos = Newpos();
        if (l == r)
        {
            siz[pos]++;
            return;
        }
        int mid = (l + r) >> 1;
        if (w <= mid)
            insert(ls[pos], w, l, mid);
        else
            insert(rs[pos], w, mid + 1, r);
        pushup(pos);
    }
    void clear(int &pos, int w, int l, int r)
    {

        if (l == r)
        {
            siz[pos]--;
            return;
        }
        int mid = (l + r) >> 1;
        if (w <= mid)
            clear(ls[pos], w, l, mid);
        else
            clear(rs[pos], w, mid + 1, r);

        pushup(pos);

        // if (siz[pos] == 0)
        // {
        //     q.push(pos);
        //     pos = 0;
        // }
    }
    int query(int ql, int qr, int pos, int l, int r)
    {
        if (siz[pos] == 0 || pos == 0)
            return 0;
        if (ql <= l && r <= qr)
        {
            return siz[pos];
        }
        int mid = (l + r) >> 1;
        int ans = 0;
        if (ql <= mid)
            ans += query(ql, qr, ls[pos], l, mid);
        if (qr > mid)
            ans += query(ql, qr, rs[pos], mid + 1, r);
        return ans;
    }
} t;
int n, k;
vector<int> g[N];
int wi[N];
int rt[N];
int dep[N], siz[N], son[N];
void dfs1(int x, int fa)
{
    dep[x] = dep[fa] + 1;
    siz[x] = 1;

    for (int to : g[x])
    {
        if (to == fa)
            continue;
        dfs1(to, x);

        siz[x] += siz[to];
        if (siz[son[x]] < siz[to])
            son[x] = to;
    }
}

int Nowson;

vector<int> p;
void dfs3(int x, int fa)
{
    p.push_back(x);
    for (int to : g[x])
    {
        if (to == fa)
            continue;
        dfs3(to, x);
    }
}
void count(int x, int fa)
{

    for (int to : g[x])
    {
        if (to == fa || to == Nowson)
            continue;
        p.clear();

        dfs3(to, x);

        for (int i : p)
        {
            int o = 2 * dep[x] + k - dep[i];
            int z = 2 * wi[x] - wi[i];

            if (z >= 0 && z <= n)
            {
                res += t.query(1, min(o, n), rt[z], 1, n);
                //cout << t.query(1, min(o, n), rt[z], 1, n) << endl;
            }
        }
        for (int i : p)
        {

            t.insert(rt[wi[i]], dep[i], 1, n);
        }
    }

    t.insert(rt[wi[x]], dep[x], 1, n);
}
void clear(int x, int fa)
{
    p.clear();
    p.push_back(x);
    for (int to : g[x])
    {
        if (to == fa)
            continue;
        dfs3(to, x);
    }
    for (int i : p)
    {
        t.clear(rt[wi[i]], dep[i], 1, n);
    }
}
void dfs2(int x, int fa, bool op)
{

    for (int to : g[x])
    {
        if (to == fa || to == son[x])
            continue;
        dfs2(to, x, 0);
    }
    if (son[x])
        dfs2(son[x], x, 1);
    Nowson = son[x];
    count(x, fa);
    Nowson = 0;
    if (!op)
        clear(x, fa);
}
int main()
{
    n = read(), k = read();
    for (int i = 1; i <= n; i++)
        wi[i] = read();
    for (int i = 2; i <= n; i++)
    {
        int x = read();
        g[x].push_back(i);
    }

    dfs1(1, 0);
    dfs2(1, 0, 0);
    cout << res * 2 << endl;
}