オイラーツアー練習第二弾。
そして、これにて一応、ツリーの辺に重みがある場合にも対応できることとなった。
問題概要
頂点 1 を根とした頂点数 N の根付き木が与えられる。初期状態では各頂点に 0 か 1 の値が割り当てられている。以下の Q 個のクエリに答えよ:
- type 1: 頂点 v を根とする部分木に含まれる頂点すべての値を反転する
- type 2: 頂点 u と頂点 v との間のパスが「０」と「１」が交互に現れるかを判定する
制約
- 2 <= N <= 105
- 1 <= Q <= 105
解法
部分木に関する問題だからやはり第一感はオイラーツアーをしたくなる。ただ今回は「部分木」に関する問題を巧みに言い換えることで、部分木に関する問題ではなくなる。
ツリーの各辺について
- 両端が同色: コスト 1
- 両端が異色: コスト 0
としてコストがついているものと考えれば
- type 1: 頂点 v とその親との間の辺の重みを反転する
- type 2: 頂点 u, v 間のパスのコストを求め、それが 0 かどうかを判定する
という風に捉え直すことができる。こうすると蟻本の例題「HouseWife Wind」と一緒である。なお、type 1 について、今回は辺の index を管理しなくてもできる。
#include <iostream> #include <vector> #include <algorithm> using namespace std; // RMQ (to find LCA) template<class VAL> struct RMQ { vector<pair<VAL, int> > dat; int SIZE_R; VAL INF = 1<<29; // to be set RMQ(int n = 110000) { init(n); } void init(int n) { SIZE_R = 1; while (SIZE_R < n) SIZE_R *= 2; dat.resize(SIZE_R * 2 - 1); for (int i = 0; i < (int)dat.size(); ++i) dat[i] = make_pair(INF, -1); } inline void set(int a, VAL v) { int k = a + SIZE_R - 1; dat[k] = make_pair(v, a); while (k > 0) { k = (k-1)/2; dat[k] = min(dat[k*2+1], dat[k*2+2]); } } inline pair<VAL,int> get(int a, int b, int k, int l, int r) { if (r <= a || b <= l) return make_pair(INF, -1); if (a <= l && r <= b) return dat[k]; else { pair<VAL,int> vl = get(a, b, k*2+1, l, (l+r)/2); pair<VAL,int> vr = get(a, b, k*2+2, (l+r)/2, r); return min(vl, vr); } } inline pair<VAL,int> get(int a, int b) { return get(a, b, 0, 0, SIZE_R); } void print() { for (int i = 0; i < SIZE_R; ++i) { VAL val = (*this)[i]; if (val < INF) cout << val; else cout << "INF"; if (i != SIZE_R-1) cout << ","; } cout << endl; } }; // BIT (to find cost of path in the tree) template<class Abel> struct BIT { vector<Abel> dat; Abel UNITY_SUM = 0; // to be set /* [1, n] */ BIT(int n = 110000) { init(n); } void init(int n) { dat.resize(n + 1); for (int i = 0; i < (int)dat.size(); ++i) dat[i] = UNITY_SUM; } /* a is 1-indexed */ inline void add(int a, Abel x) { for (int i = a; i < (int)dat.size(); i += i & -i) dat[i] = dat[i] + x; } /* [1, a], a is 1-indexed */ inline Abel sum(int a) { Abel res = UNITY_SUM; for (int i = a; i > 0; i -= i & -i) res = res + dat[i]; return res; } /* [a, b), a and b are 1-indexed */ inline Abel sum(int a, int b) { return sum(b - 1) - sum(a - 1); } /* a is 1-indexed */ inline Abel operator [] (int a) { return sum(a, a+1); } void print() { for (int i = 1; i < (int)dat.size(); ++i) cout << sum(i, i + 1) << ","; cout << endl; } }; // Graph template<class VAL> struct Edge { int idx, from, to; VAL cost; Edge(int idx_, int from_, int to_, VAL cost_) : idx(idx_), from(from_), to(to_), cost(cost_) {} friend ostream& operator << (ostream& s, const Edge<VAL>& e) { return s << e.from << "->" << e.to << '(' << e.cost << ')'; } }; template<class VAL> struct Graph { int iter = 0; vector<vector<Edge<VAL> > > list; Graph(int n = 110000) : iter(0) { init(n); } void init(int n = 110000) { iter = 0; list.clear(); list.resize(n); } inline vector<Edge<VAL> > operator [] (int i) {return list[i];} size_t size() const { return list.size(); } void connect(int f, int t, VAL c) { list[f].push_back(Edge<VAL>(iter, f, t, c)); list[t].push_back(Edge<VAL>(iter, t, f, c)); ++iter; } friend ostream& operator << (ostream& s, const Graph& G) { s << endl; for (int i = 0; i < (int)G.size(); ++i) {s << i << " : " << G.list[i] << endl;} return s; } }; // Euler Tour template<class VAL> struct EulerTour { // main results Graph<VAL> tree; vector<int> depth; vector<int> node; // the node-number of i-th element of Euler-tour vector<int> vf, ve; // the index of Euler-tour of node v vector<int> eid; // the index of edge e (i*2 + (0: dir to leaf, 1: dir to root)) // sub results RMQ<int> rmq; // depth (to find LCA) BIT<VAL> bit; // to calc sum of edges EulerTour(const Graph<VAL> &tree_) { init(tree_); } void init(const Graph<VAL> &tree_) { tree = tree_; int V = (int)tree.size(); depth.resize(V*2-1); node.resize(V*2-1); vf.resize(V); ve.resize(V); eid.resize((V-1)*2); rmq.init(V*2-1); bit.init((V-1)*2); int k = 0; dfs(0, -1, 0, k); for (int i = 0; i < V*2-1; ++i) rmq.set(i, depth[i]); } void dfs(int v, int par, int dep, int &ord) { node[ord] = v; depth[ord] = dep; vf[v] = ve[v] = ord; ++ord; for (auto e : tree[v]) { if (e.to == par) continue; bit.add(ord, e.cost); eid[e.idx*2] = ord; dfs(e.to, v, dep+1, ord); node[ord] = v; depth[ord] = dep; ve[v] = ord; eid[e.idx*2+1] = ord; bit.add(ord, -e.cost); // minus cost for opposite direction ++ord; } } inline int LCA(int u, int v) { int a = vf[u], b = vf[v]; if (a > b) swap(a, b); return node[rmq.get(a, b+1).second]; } inline void add(int index_edge, VAL v) { bit.add(eid[index_edge*2], v); bit.add(eid[index_edge*2+1], -v); } inline VAL sum(int u, int v) { int lca = LCA(u, v); return bit.sum(vf[u]) + bit.sum(vf[v]) - bit.sum(vf[lca])*2; } //void print() { COUT(node); COUT(vf); COUT(eid); bit.print(); } }; int N, Q; vector<int> P, C; Graph<int> tree; int main() { cin >> N; tree.init(N); P.resize(N-1); C.resize(N); for (int i = 0; i < N-1; ++i) scanf("%d", &P[i]), --P[i]; for (int i = 0; i < N; ++i) scanf("%d", &C[i]); for (int i = 0; i < N-1; ++i) { int weights = 0; if (C[i+1] == C[P[i]]) weights = 1; tree.connect(i+1, P[i], weights); } EulerTour<int> et(tree); cin >> Q; for (int q = 0; q < Q; ++q) { int type; cin >> type; if (type == 1) { int v; scanf("%d", &v); --v; if (v == 0) continue; int cur = et.bit[et.vf[v]]; int add; if (cur == 0) add = 1; else add = -1; et.bit.add(et.vf[v], add); et.bit.add(et.ve[v]+1, -add); } else { int u, v; scanf("%d %d", &u, &v); --u, --v; int res = et.sum(u, v); if (res == 0) cout << "YES" << endl; else cout << "NO" << endl; } } }
