#87470 (C++11) No.363 門松サイクル

• 問題: No.363 門松サイクル
• ユーザー: anta
• 提出日時: 2016-04-17 23:23:17
• 言語: C++11 (gcc 11.4.0)
• 結果: AC
• 実行時間: 125 ms / 4,000 ms
• コード長: 9,427 bytes
• コンパイル時間: 1,282 ms
• コンパイル使用メモリ: 101,156 KB
• 実行使用メモリ: 18,696 KB
• 最終ジャッジ日時: 2024-10-04 10:55:57
• 合計ジャッジ時間: 4,212 ms
• ジャッジサーバーID （参考情報）: judge2 / judge4

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 6,816 KB
testcase_01	AC	1 ms; 6,820 KB
testcase_02	AC	2 ms; 6,820 KB
testcase_03	AC	2 ms; 6,816 KB
testcase_04	AC	2 ms; 6,816 KB
testcase_05	AC	2 ms; 6,816 KB
testcase_06	AC	3 ms; 6,816 KB
testcase_07	AC	2 ms; 6,820 KB
testcase_08	AC	3 ms; 6,820 KB
testcase_09	AC	43 ms; 9,344 KB
testcase_10	AC	48 ms; 9,856 KB
testcase_11	AC	93 ms; 15,968 KB
testcase_12	AC	106 ms; 15,276 KB
testcase_13	AC	111 ms; 15,552 KB
testcase_14	AC	90 ms; 15,236 KB
testcase_15	AC	84 ms; 15,628 KB
testcase_16	AC	60 ms; 11,512 KB
testcase_17	AC	103 ms; 17,204 KB
testcase_18	AC	122 ms; 16,756 KB
testcase_19	AC	103 ms; 18,696 KB
testcase_20	AC	125 ms; 16,652 KB
testcase_21	AC	119 ms; 17,336 KB
testcase_22	AC	117 ms; 16,184 KB
testcase_23	AC	86 ms; 15,140 KB
testcase_24	AC	2 ms; 6,816 KB
testcase_25	AC	1 ms; 6,820 KB
testcase_26	AC	73 ms; 17,328 KB
testcase_27	AC	76 ms; 17,208 KB
testcase_28	AC	92 ms; 17,600 KB

コンパイルメッセージ

main.cpp: In function ‘int main()’:
main.cpp:268:30: warning: ignoring return value of ‘int scanf(const char*, ...)’ declared with attribute ‘warn_unused_result’ [-Wunused-result]
  268 |                         scanf("%d", &A[i]);
      |                         ~~~~~^~~~~~~~~~~~~
main.cpp:272:30: warning: ignoring return value of ‘int scanf(const char*, ...)’ declared with attribute ‘warn_unused_result’ [-Wunused-result]
  272 |                         scanf("%d%d", &u, &v), -- u, -- v;
      |                         ~~~~~^~~~~~~~~~~~~~~~
main.cpp:293:22: warning: ignoring return value of ‘int scanf(const char*, ...)’ declared with attribute ‘warn_unused_result’ [-Wunused-result]
  293 |                 scanf("%d", &Q);
      |                 ~~~~~^~~~~~~~~~
main.cpp:296:30: warning: ignoring return value of ‘int scanf(const char*, ...)’ declared with attribute ‘warn_unused_result’ [-Wunused-result]
  296 |                         scanf("%d%d", &u, &v), -- u, -- v;
      |                         ~~~~~^~~~~~~~~~~~~~~~

ソースコード

#include <string>
#include <vector>
#include <algorithm>
#include <numeric>
#include <set>
#include <map>
#include <queue>
#include <iostream>
#include <sstream>
#include <cstdio>
#include <cmath>
#include <ctime>
#include <cstring>
#include <cctype>
#include <cassert>
#include <limits>
#include <functional>
#define rep(i,n) for(int (i)=0;(i)<(int)(n);++(i))
#define rer(i,l,u) for(int (i)=(int)(l);(i)<=(int)(u);++(i))
#define reu(i,l,u) for(int (i)=(int)(l);(i)<(int)(u);++(i))
#if defined(_MSC_VER) || __cplusplus > 199711L
#define aut(r,v) auto r = (v)
#else
#define aut(r,v) __typeof(v) r = (v)
#endif
#define each(it,o) for(aut(it, (o).begin()); it != (o).end(); ++ it)
#define all(o) (o).begin(), (o).end()
#define pb(x) push_back(x)
#define mp(x,y) make_pair((x),(y))
#define mset(m,v) memset(m,v,sizeof(m))
#define INF 0x3f3f3f3f
#define INFL 0x3f3f3f3f3f3f3f3fLL
using namespace std;
typedef vector<int> vi; typedef pair<int, int> pii; typedef vector<pair<int, int> > vpii; typedef long long ll;
template<typename T, typename U> inline void amin(T &x, U y) { if(y < x) x = y; }
template<typename T, typename U> inline void amax(T &x, U y) { if(x < y) x = y; }

class SchieberVishkinLCA {
public:
	typedef unsigned Word;
	typedef int Vertex;
private:

	static inline Word lowestOneBit(Word v) {
		return v & (~v + 1);
	}
	static inline Word highestOneBitMask(Word v) {
		v |= v >> 1; v |= v >> 2; v |= v >> 4; v |= v >> 8; v |= v >> 16;
		return v >> 1;
	}

	std::vector<Word> indices;			//Vertex -> index
	std::vector<Word> maxHIndices;		//Vertex -> index
	std::vector<Word> ancestorHeights;	//Vertex -> Word
	std::vector<Vertex> pathParents;	//index-1 -> Vertex
public:

	void build(const std::vector<Vertex> &preorder, const std::vector<Vertex> &parents, Vertex root) {
		Vertex N = static_cast<Vertex>(preorder.size());

		ancestorHeights.resize(N);
		maxHIndices.resize(N);
		indices.resize(N);
		pathParents.resize(N);

		for(Vertex ix = 0; ix < N; ++ ix)
			indices[preorder[ix]] = ix + 1;

		for(Vertex i = 0; i < N; ++ i)
			maxHIndices[i] = indices[i];
		for(Vertex ix = N - 1; ix > 0; -- ix) {
			Vertex v = preorder[ix], parent = parents[v];
			if(lowestOneBit(maxHIndices[parent]) < lowestOneBit(maxHIndices[v]))
				maxHIndices[parent] = maxHIndices[v];
		}

		ancestorHeights[root] = 0;
		for(Vertex ix = 1; ix < N; ++ ix) {
			Vertex v = preorder[ix], parent = parents[v];
			ancestorHeights[v] = ancestorHeights[parent] | lowestOneBit(maxHIndices[v]);
		}

		pathParents[0] = root;
		for(Vertex ix = 1; ix < N; ++ ix) {
			Vertex v = preorder[ix], parent = parents[v];
			if(maxHIndices[v] != maxHIndices[parent])
				pathParents[indices[v] - 1] = parent;
			else
				pathParents[indices[v] - 1] = pathParents[indices[parent] - 1];
		}
	}

	Vertex query(Vertex v, Vertex u) const {
		Word Iv = maxHIndices[v], Iu = maxHIndices[u];
		Word hIv = lowestOneBit(Iv), hIu = lowestOneBit(Iu);
		Word hbMask = highestOneBitMask((Iv ^ Iu) | hIv | hIu);
		Word j = lowestOneBit(ancestorHeights[v] & ancestorHeights[u] & ~hbMask);
		Vertex x, y;
		if(j == hIv) x = v;
		else {
			Word kMask = highestOneBitMask(ancestorHeights[v] & (j - 1));
			x = pathParents[(indices[v] & ~kMask | (kMask + 1)) - 1];
		}
		if(j == hIu) y = u;
		else {
			Word kMask = highestOneBitMask(ancestorHeights[u] & (j - 1));
			y = pathParents[(indices[u] & ~kMask | (kMask + 1)) - 1];
		}
		return indices[x] < indices[y] ? x : y;
	}
};

bool isok(int a, int b, int c) {
	return a != b && a != c && b != c && (min({ a, b, c }) == b || max({ a, b, c }) == b);
}

vector<int> t_parent;
vi t_ord;

void tree_getorder(const vector<vi> &g, int root) {
	int n = g.size();
	t_parent.assign(n, -1);
	t_ord.clear();

	vector<int> stk; stk.push_back(root);
	while(!stk.empty()) {
		int i = stk.back(); stk.pop_back();
		t_ord.push_back(i);
		for(int j = (int)g[i].size() - 1; j >= 0; j --) {
			int c = g[i][j];
			if(t_parent[c] == -1 && c != root)
				stk.push_back(c);
			else
				t_parent[i] = c;
		}
	}
}

struct HeavyLightDecomposition {
	vector<int> colors, positions;	//Vertex -> Color, Vertex -> Offset
	vector<int> lengths, parents, branches;	//Color -> Int, Color -> Color, Color -> Offset
	vector<int> parentnodes, depths;	//Vertex -> Vertex, Vertex -> Int
										//vector<FenwickTree>とかを避けて1次元にしたい時に使う
										//sortednodesの[lefts[v], rights[v])はvのsubtreeとなっている
	vector<int> sortednodes, offsets;	//Index -> Vertex, Color -> Index
	vector<int> lefts, rights;	//Vertex -> Index

	struct BuildDFSState {
		int i, len, parent;
		BuildDFSState() {}
		BuildDFSState(int i_, int l, int p) : i(i_), len(l), parent(p) {}
	};

	//両方の辺があってもいいし、親から子への辺だけでもよい
	void build(const vector<vi> &g, int root) {
		int n = g.size();

		colors.assign(n, -1); positions.assign(n, -1);
		lengths.clear(); parents.clear(); branches.clear();
		parentnodes.assign(n, -1); depths.assign(n, -1);

		sortednodes.clear(); offsets.clear();
		lefts.assign(n, -1); rights.assign(n, -1);

		vector<int> subtreesizes;
		measure(g, root, subtreesizes);

		typedef BuildDFSState State;
		depths[root] = 0;
		vector<State> s;
		s.push_back(State(root, 0, -1));
		while(!s.empty()) {
			State t = s.back(); s.pop_back();
			int i = t.i, len = t.len;
			int index = sortednodes.size();
			int color = lengths.size();

			if(t.parent == -3) {
				rights[i] = index;
				continue;
			}

			if(t.parent != -2) {
				assert(parents.size() == color);
				parents.push_back(t.parent);
				branches.push_back(len);
				offsets.push_back(index);
				len = 0;
			}
			colors[i] = color;
			positions[i] = len;

			lefts[i] = index;
			sortednodes.push_back(i);

			int maxsize = -1, maxj = -1;
			each(j, g[i]) if(colors[*j] == -1) {
				if(maxsize < subtreesizes[*j]) {
					maxsize = subtreesizes[*j];
					maxj = *j;
				}
				parentnodes[*j] = i;
				depths[*j] = depths[i] + 1;
			}
			s.push_back(State(i, -1, -3));
			if(maxj == -1) {
				lengths.push_back(len + 1);
			} else {
				each(j, g[i]) if(colors[*j] == -1 && *j != maxj)
					s.push_back(State(*j, len, color));
				s.push_back(State(maxj, len + 1, -2));
			}
		}
	}

	void get(int v, int &c, int &p) const {
		c = colors[v]; p = positions[v];
	}
	bool go_up(int &c, int &p) const {
		p = branches[c]; c = parents[c];
		return c != -1;
	}

	inline const int *nodesBegin(int c) const { return &sortednodes[0] + offsets[c]; }
	inline const int *nodesEnd(int c) const { return &sortednodes[0] + (c + 1 == offsets.size() ? sortednodes.size() : offsets[c + 1]); }

private:
	void measure(const vector<vi> &g, int root, vector<int> &out_subtreesizes) const {
		out_subtreesizes.assign(g.size(), -1);
		vector<int> s;
		s.push_back(root);
		while(!s.empty()) {
			int i = s.back(); s.pop_back();
			if(out_subtreesizes[i] == -2) {
				int s = 1;
				each(j, g[i]) if(out_subtreesizes[*j] != -2)
					s += out_subtreesizes[*j];
				out_subtreesizes[i] = s;
			} else {
				s.push_back(i);
				each(j, g[i]) if(out_subtreesizes[*j] == -1)
					s.push_back(*j);
				out_subtreesizes[i] = -2;
			}
		}
	}
};

int level_ancestor(const HeavyLightDecomposition &hld, int x, int d) {
	int c, p;
	const int *head;
	if(d > hld.depths[x]) return -1;
	hld.get(x, c, p);
	while(d < hld.depths[*(head = hld.nodesBegin(c))])
		hld.go_up(c, p);
	return head[d - hld.depths[*head]];
}
int get_ancestor_child(const HeavyLightDecomposition &hld, int x, int y) {
	return level_ancestor(hld, y, hld.depths[x] + 1);
}


int main() {
	int N;
	while(~scanf("%d", &N)) {
		vector<int> A(N);
		for(int i = 0; i < N; ++ i)
			scanf("%d", &A[i]);
		vector<vector<int> > g(N);
		for(int i = 0; i < N - 1; ++ i) {
			int u, v;
			scanf("%d%d", &u, &v), -- u, -- v;
			g[u].push_back(v);
			g[v].push_back(u);
		}
		tree_getorder(g, 0);
		vector<bool> parok(N);
		for(int ix = 0; ix < (int)t_ord.size(); ++ ix) {
			int i = t_ord[ix], p = t_parent[i];
			parok[i] = p == -1 || p == 0 || isok(A[t_parent[p]], A[p], A[i]);
		}
		vector<int> okcnt(N), depth(N);
		okcnt[0] = 1;
		for(int ix = 1; ix < (int)t_ord.size(); ++ ix) {
			int i = t_ord[ix], p = t_parent[i];
			okcnt[i] = okcnt[p] + parok[i];
			depth[i] = depth[p] + 1;
		}
		SchieberVishkinLCA lca;
		lca.build(t_ord, t_parent, 0);
		HeavyLightDecomposition hld; hld.build(g, 0);
		int Q;
		scanf("%d", &Q);
		rep(ii, Q) {
			int u; int v;
			scanf("%d%d", &u, &v), -- u, -- v;
			int w = lca.query(u, v);
			int len = depth[u] + depth[v] - depth[w] * 2 + 1;
			bool ans;
			if(len <= 2) {
				ans = false;
			} else {
				ans = true;
				int num = 0;
				//wが真ん中
				if(w != u && w != v) {
					int a = level_ancestor(hld, u, hld.depths[w] + 1);
					int b = level_ancestor(hld, v, hld.depths[w] + 1);
					ans &= isok(A[a], A[w], A[b]);
					num += 1;
				}
				//u～w, v～w
				rep(uv, 2) {
					int t = uv == 0 ? u : v;
					if(t == w) continue;
					int a = level_ancestor(hld, t, hld.depths[w] + 1);
					ans &= okcnt[t] - okcnt[a] == depth[t] - depth[a];
					num += depth[t] - depth[a];
				}
				//追加された辺を含むもの
				rep(uv, 2) {
					int x = uv == 0 ? u : v, y = uv == 0 ? v : u;
					int c;
					if(y == w) {
						c = level_ancestor(hld, x, hld.depths[w] + 1);
					} else {
						c = t_parent[y];
					}
					ans &= isok(A[x], A[y], A[c]);
					num += 1;
				}
				assert(num == len);
			}
			puts(ans ? "YES" : "NO");
		}
	}
	return 0;
}