题意就是说要数有多少组 $(u, v, w)$，使得 $u, v, w$ 的类型分别为 $0, 1, 2$，且 $v$ 在 $u, w$ 的最短路径上。那么，我们就对这种链进行计数。

考虑树剖然后 dp。设 $f_{i, 0/1/2/3/4}$ 表示只考虑 $i$ 的子树，

0. 当前有几条只匹配了 $0$ 类点的链。
1. 当前有几条只匹配了 $0, 1$ 类点的链。
2. 当前有几条只匹配了 $2$ 类点的链。
3. 当前有几条只匹配了 $2, 1$ 类点的链。
4. 当前有几条匹配完成的链。

答案就是 $f_{1, 4}$。

设 $u$ 的重儿子是 $s$，那么我们希望写出一个矩阵 $G_u$，使得 $\begin{bmatrix}f_{s, 0} & \dots & f_{s, 4}\end{bmatrix} \times G_u = \begin{bmatrix}f_{u, 0} & \dots & f_{u, 4}\end{bmatrix}$（当然后面会发现要补一项 $1$）。

设 $g_{u, 0/1/2/3/4}$ 表示 $u$ 的所有轻儿子的信息并，下标和 $f$ 的定义相同。那么，在加入轻儿子 $v$ 时：

$$\begin{cases} g'_{u, 0} \gets g_{u, 0} + f_{v, 0}\\ g'_{u, 1} \gets g_{u, 1} + f_{v, 1}\\ g'_{u, 2} \gets g_{u, 2} + f_{v, 2}\\ g'_{u, 3} \gets g_{u, 3} + f_{v, 3}\\ g'_{u, 4} \gets g_{u, 4} + f_{v, 4} + \sum_{i=0}^3 g_{u, i}f_{v, 3-i} \end{cases}$$

这样其实还漏了一种情况，就是当 $u$ 类型为 $1$ 时，可能有两个来自不同子树的 $0, 2$ 类点，它们可能在这里合并为一条完整的链。因此补充定义 $g_{u, 5}$ 表示这种情况。可以得到：

$$g'_{u, 5} \gets g_{u, 5} + g_{u, 0}f_{v, 2} + g_{u, 2}f_{v, 0}$$

那么，我们可以把 $g_u$ 转成转移矩阵 $G_u$。如果 $u$ 为 $0$（空着的位置为 $0$）：

$$\def\mat#1{\begin{bmatrix}#1\end{bmatrix}}\\ \mat{f_{v, 0} & f_{v, 1} & f_{v, 2} & f_{v, 3} & f_{v, 4} & 1} \times \mat{ 1 & & & & g_{u, 3} & \\ & 1 & & & g_{u, 2} & \\ & & 1 & & g_{u, 1} & \\ & & & 1 & g_{u, 0} + 1 & \\ & & & & 1 & \\ g_{u, 0} + 1 & g_{u, 1} & g_{u, 2} & g_{u, 3} & g_{u, 3} + g_{u, 4} & 1 \\ } = \mat{f_{u, 0} & f_{u, 1} & f_{u, 2} & f_{u, 3} & f_{u, 4} & 1}$$

如果 $u$ 为 $2$，这也是类似的：

$$\def\mat#1{\begin{bmatrix}#1\end{bmatrix}}\\ \mat{f_{v, 0} & f_{v, 1} & f_{v, 2} & f_{v, 3} & f_{v, 4} & 1} \times \mat{ 1 & & & & g_{u, 3} & \\ & 1 & & & g_{u, 2} + 1 & \\ & & 1 & & g_{u, 1} & \\ & & & 1 & g_{u, 0} & \\ & & & & 1 & \\ g_{u, 0} & g_{u, 1} & g_{u, 2} + 1 & g_{u, 3} & g_{u, 1} + g_{u, 4} & 1 \\ } = \mat{f_{u, 0} & f_{u, 1} & f_{u, 2} & f_{u, 3} & f_{u, 4} & 1}$$

如果 $u$ 为 $1$，那么会复杂一点：

$$\def\mat#1{\begin{bmatrix}#1\end{bmatrix}}\\ \mat{f_{v, 0} & f_{v, 1} & f_{v, 2} & f_{v, 3} & f_{v, 4} & 1} \times \mat{ 1 & 1 & & & g_{u, 2} + g_{u, 3} & \\ & 1 & & & g_{u, 2} & \\ & & 1 & 1 & g_{u, 0} + g_{u, 1} & \\ & & & 1 & g_{u, 0} & \\ & & & & 1 & \\ g_{u, 0} & g_{u, 0} + g_{u, 1} & g_{u, 2} & g_{u, 2} + g_{u, 3} & g_{u, 4} + g_{u, 5} & 1 \\ } = \mat{f_{u, 0} & f_{u, 1} & f_{u, 2} & f_{u, 3} & f_{u, 4} & 1}$$

直接树剖 + sgt 维护动态 dp 可以做到 $O(n\log^2n \times \delta^3)$，其中 $\delta = 6$ 为矩阵大小。然后就可以过了。

用全局平衡二叉树可以少一个 $\log$，但是我没写。

:::success[代码]

// #include "joitour.h"

#include <algorithm>
#include <vector>
using namespace std;

namespace P10437
{

#define MAXN 200005

using ll = long long;

struct Mat
{
    ll m[6][6];
    inline ll *operator[](size_t v)
    {
        return m[v];
    }
    inline const ll *operator[](size_t v) const
    {
        return m[v];
    }

    inline Mat operator*(const Mat &b) const
    {
        Mat c{};
        for (int i = 0; i < 6; i++)
        {
            for (int j = 0; j < 6; j++)
            {
                for (int k = 0; k < 6; k++)
                {
                    c[i][k] += m[i][j] * b[j][k];
                }
            }
        }
        return c;
    }
} I;

struct Vec
{
    ll m[6];
    inline ll &operator[](size_t v)
    {
        return m[v];
    }
    inline ll operator[](size_t v) const
    {
        return m[v];
    }

    inline Vec &operator+=(const Vec &b)
    {
        for (int i = 0; i < 4; i++)
        {
            m[4] += m[i] * b[3 - i];
        }
        m[5] += m[0] * b[2] + m[2] * b[0];
        for (int i = 0; i < 5; i++)
        {
            m[i] += b[i];
        }
        return *this;
    }

    inline Vec &operator-=(const Vec &b)
    {
        for (int i = 0; i < 5; i++)
        {
            m[i] -= b[i];
        }
        m[5] -= m[0] * b[2] + m[2] * b[0];
        for (int i = 0; i < 4; i++)
        {
            m[4] -= m[i] * b[3 - i];
        }
        return *this;
    }

    inline static Vec from(const Mat &mat)
    {
        Vec v{};
        for (int i = 0; i < 5; i++)
        {
            v[i] += mat[5][i];
        }
        return v;
    }
    inline static Mat to(const Vec &vec, int t)
    {
        Mat mat = I;
        for (int i = 0; i < 4; i++)
        {
            mat[5][i] += vec[i];
            mat[3 - i][4] += vec[i];
        }
        mat[5][4] += vec[4];
        if (t == 0)
        {
            mat[3][4]++;
            mat[5][0]++;
            mat[5][4] += vec[3];
        }
        else if (t == 2)
        {
            mat[1][4]++;
            mat[5][2]++;
            mat[5][4] += vec[1];
        }
        else
        {
            mat[0][1]++;
            mat[2][3]++;
            mat[0][4] += vec[2];
            mat[2][4] += vec[0];
            mat[5][1] += vec[0];
            mat[5][3] += vec[2];
            mat[5][4] += vec[5];
        }
        return mat;
    }
};

int n;

int col[MAXN];

vector<int> e[MAXN];

int fa[MAXN], son[MAXN], siz[MAXN];
int dfn[MAXN], top[MAXN], tai[MAXN], dfc = 0;
int idn[MAXN];

void dfs1(int p, int f)
{
    fa[p] = f;
    siz[p] = 1;
    for (int u : e[p])
    {
        if (u == f)
        {
            continue;
        }
        dfs1(u, p);
        siz[p] += siz[u];
        if (siz[u] > siz[son[p]])
        {
            son[p] = u;
        }
    }
}

Mat t[MAXN * 3];
int P;
#define ls (p << 1)
#define rs (p << 1 | 1)

inline void build()
{
    P = 1;
    while (P <= n + 1)
    {
        P <<= 1;
    }
    t[P] = I;
    for (int i = 1; i <= n; i++)
    {
        t[i + P] = I;
    }
    for (int p = P - 1; p; p--)
    {
        if (rs <= P + n + 1)
        {
            t[p] = t[rs] * t[ls];
        }
    }
}

inline Mat query(int l, int r)
{
    l += P - 1, r += P + 1;
    Mat ql = I, qr = I;
    while (l ^ r ^ 1)
    {
        if (~l & 1)
        {
            ql = t[l ^ 1] * ql;
        }
        if (r & 1)
        {
            qr = qr * t[r ^ 1];
        }
        l >>= 1, r >>= 1;
    }
    return qr * ql;
}

inline void update(int p, const Mat &v)
{
    t[p += P] = v;
    for (p >>= 1; p; p >>= 1)
    {
        if (rs <= P + n + 1)
        {
            t[p] = t[rs] * t[ls];
        }
    }
}

Vec G[MAXN], lu[MAXN];

inline Vec calc(int p)
{
    return Vec::from(query(dfn[p], dfn[tai[p]]));
}

void dfs2(int p, int t)
{
    top[p] = t;
    tai[t] = p;
    dfn[p] = ++dfc;
    idn[dfc] = p;
    if (!son[p])
    {
        return update(dfn[p], Vec::to(G[p], col[p]));
    }
    dfs2(son[p], t);
    for (int u : e[p])
    {
        if (u == fa[p] || u == son[p])
        {
            continue;
        }
        dfs2(u, u);
        G[p] += lu[u] = calc(u);
    }
    update(dfn[p], Vec::to(G[p], col[p]));
}

}

void init(int _N, vector<int> _F, vector<int> _U, vector<int> _V, int)
{
    using namespace P10437;
    for (int i = 0; i < 6; i++)
    {
        I[i][i] = 1;
    }
    n = _N;
    copy(_F.begin(), _F.end(), col + 1);
    for (int i = 0; i < n - 1; i++)
    {
        e[++_U[i]].push_back(++_V[i]);
        e[_V[i]].push_back(_U[i]);
    }
    dfs1(1, 1);
    build();
    dfs2(1, 1);
}

void change(int u, int v)
{
    using namespace P10437;
    col[++u] = v;
    update(dfn[u], Vec::to(G[u], col[u]));
    for (int p = top[u]; p != 1; p = top[fa[p]])
    {
        G[fa[p]] -= lu[p];
        G[fa[p]] += lu[p] = calc(p);
        update(dfn[fa[p]], Vec::to(G[fa[p]], col[fa[p]]));
    }
}

long long num_tours()
{
    using namespace P10437;
    return calc(1)[4];
}

:::