#include "joitour.h"
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#define N 200000
int n, color[N], hd[N], e[N * 2], Ln[N * 2], ii, sz[N], hld[N], tail[N], nfd[N], dfn[N], dfc, fa[N], lv[N], n_[N], childcount[N], id[N], m_[N];
struct path {
long long ans, J, O, I, OI, JO, IO, OJ, I_, J_;
path() {
ans = J = O = I = OI = JO = IO = OJ = I_ = J_ = 0;
}
friend path operator+(path l, path r) { /* compress */
path c;
c.ans = l.ans + l.JO * r.I + l.IO * r.J + r.ans + l.I_ * r.OJ + l.J_ * r.OI;
c.J = l.J + r.J;
c.O = l.O + r.O;
c.I = l.I + r.I;
c.JO = l.J * r.O + l.JO * r.JO;
c.IO = l.I * r.O + l.IO + r.IO;
c.OI = l.O * r.I + l.OI + r.OI;
c.OJ = l.O * r.J + l.OJ + r.OJ;
if (r.O) {
c.I_ = r.I_;
c.J_ = r.J_;
} else {
c.I_ = l.I_ + r.I_;
c.J_ = l.J_ + r.J_;
}
return c;
}
void print() {
return;
printf(" PATH {ans = %lld, J = %lld, O = %lld, I = %lld, OI = %lld, JO = %lld, IO = %lld, OJ = %lld, I_ = %lld, J_ = %lld}\n", ans, J, O, I, OI, JO, IO, OJ, I_, J_);
}
};
struct point {
long long ans, J, O, I, JO, IO, /* cross subtree */ IJ, I_, J_;
point() {
ans = J = O = I = JO = IO = IJ = I_ = J_ = 0;
}
friend point operator+(point l, point r) { /* rake */
point c;
c.ans = l.ans + r.ans;
c.J = l.J + r.J;
c.O = l.O + r.O;
c.I = l.I + r.I;
c.JO = l.JO + r.JO;
c.IO = l.IO + r.IO;
c.IJ = l.IJ + r.IJ + l.J * r.I + l.I * r.J;
c.I_ = l.I_ + r.I_;
c.J_ = l.J_ + r.J_;
return c;
}
};
path *tt[N];
point *yy[N];
point Addedge(path x) {
point p;
p.ans = x.ans;
p.J = x.J;
p.O = x.O;
p.I = x.I;
p.JO = x.JO;
p.IO = x.IO;
p.IJ = 0;
p.I_ = x.I_;
p.J_ = x.J_;
return p;
}
path Addvertex(point p, char col) {
path x;
x.ans = p.ans;
x.I_ = p.I_;
x.J_ = p.J_;
if (col == 0) {
++x.J_;
x.ans += p.IO;
} else if (col == 1) {
x.ans += p.IJ;
x.I_ = x.J_ = 0;
} else if (col == 2) {
++x.I_;
x.ans += p.JO;
}
x.J = p.J;
x.O = p.O;
x.I = p.I;
x.JO = p.JO;
x.IO = p.IO;
return x;
}
path Vertex(int col) {
path x;
if (col == 0) {
x.J = x.J_ = 1;
} else if (col == 1) {
x.O = 1;
} else if (col == 2) {
x.I = x.I_ = 1;
}
return x;
}
void add(int u, int v) {
++ii;
Ln[ii] = hd[u];
e[ii] = v;
hd[u] = ii;
}
void dfs1(int u) {
sz[u] = 1;
for (int v, j = hd[u]; j; j = Ln[j]) {
v = e[j];
if (v == fa[u])
continue;
fa[v] = u;
dfs1(v);
sz[u] += sz[v];
if (e[hd[u]] == fa[u] || sz[e[hd[u]]] < sz[v])
e[j] = e[hd[u]], e[hd[u]] = v;
++childcount[u];
}
}
void dfs2(int u) {
nfd[dfn[u] = dfc++] = u;
tail[hld[u]] = u;
int ch = 0;
for (int v, j = hd[u]; j; j = Ln[j]) {
v = e[j];
if (v == fa[u])
continue;
id[v] = ch++;
hld[v] = (j != hd[u]) ? v: hld[u];
lv[v] = lv[u] + 1;
dfs2(v);
}
}
#define md ((l + r) / 2)
void build(path *t, int v, int l, int r) {
if (l == r) {
t[v] = path();
if (color[v] == 0)
t[v].J = 1;
else if (color[v] == 1)
t[v].O = 1;
else if (color[v] == 2)
t[v].I = 1;
} else {
build(t, v << 1, l, md);
build(t, v << 1 | 1, md + 1, r);
t[v] = t[v << 1] + t[v << 1 | 1];
}
}
void pul(path *t, int v) {
t[v] = t[v << 1 | 1] + t[v << 1];
}
void pull(point *t, int v) {
t[v] = t[v << 1] + t[v << 1 | 1];
}
void re_lp(int u) {
int p = fa[u];
if (p == u)
return;
yy[p][id[u] + m_[p]] = Addedge(tt[u][1]);
for (int j = id[u] + m_[p]; j >>= 1; )
pull(yy[p], j);
tt[hld[p]][dfn[p] - dfn[hld[p]] + n_[hld[p]]] = Addvertex(yy[p][1], color[p]);
for (int j = dfn[p] - dfn[hld[p]] + n_[hld[p]]; j >>= 1; )
pul(tt[hld[p]], j);
}
void dfs3(int u) {
for (int v, j = hd[u]; j; j = Ln[j]) {
v = e[j];
if (v == fa[u])
continue;
dfs3(v);
}
if (fa[u] != u && dfn[fa[u]] + 1 != dfn[u])
re_lp(u);
}
void init(int NN, std::vector<int> F, std::vector<int> U, std::vector<int> V, int Q) {
n = NN;
for (int i = 0; i + 1 < n; ++i)
add(U[i], V[i]), add(V[i], U[i]);
memcpy(color, F.data(), sizeof(int) * n);
fa[0] = 0;
dfs1(0);
hld[0] = 0;
lv[0] = 0;
dfs2(0);
/* build compress tree */
for (int i = 0; i < n; ++i) {
tail[i] = tail[hld[i]];
if (hld[i] != i)
continue;
int len = dfn[tail[i]] - dfn[i] + 1;
for (n_[i] = 1; n_[i] < len; n_[i] <<= 1);
tt[i] = new path[n_[i] << 1];
for (int j = dfn[i]; j <= dfn[tail[i]]; ++j)
tt[i][j + n_[i] - dfn[i]] = Vertex(color[nfd[j]]);
for (int j = n_[i] - 1; j >= 1; --j)
pul(tt[i], j);
}
/* build rake tree */
for (int i = 0; i < n; ++i) {
if (! childcount[i])
continue;
for (m_[i] = 1; m_[i] < childcount[i]; m_[i] <<= 1);
yy[i] = new point[m_[i] << 1];
}
dfs3(0);
////for (int i = 0; i < n; ++i) printf(" [%d] - id = %d, dfn = %d, hld = %d, tail = %d\n", i, id[i], dfn[i], hld[i], tail[i]);
}
void change(int X, int Y) {
color[X] = Y;
for (int p = X; p; p = fa[hld[p]]) {
tt[hld[p]][dfn[p] - dfn[hld[p]] + n_[hld[p]]] = childcount[p]?
Addvertex(yy[p][1], color[p]): Vertex(color[p]);
for (int j = dfn[p] - dfn[hld[p]] + n_[hld[p]]; j >>= 1; )
pul(tt[hld[p]], j);
re_lp(hld[p]);
if (hld[p] == 0)
break;
}
}
long long num_tours() {
tt[0][3].print();
tt[0][2].print();
tt[0][1].print();
path x = tt[0][1];
return x.ans;
}
