Atcoder解いてみた AtCoder Beginner Contest E - Count Descendants

AtCoder Beginner Contest E - Count Descendants

問題はこちら

回答

素直にやると、O(NQ)となってしまう
図のように、DFSし、各点に対して、行きかけ順、帰りかけ順を求める
そして、深さ毎に各頂点の行きかけ順を入れていく
深さ毎の行きかけ順の情報は昇順になるようにする
そして、クエリ毎に以下の2つの処理を行い、結果1-結果2を行うことで、答えが求まる

1. 深さ毎に行きがけ順番号を昇順に格納した配列に対して、帰りがけ順の配列の要素uの値(クエリ)を用いてlower_bounds
2. 深さ毎に行きがけ順番号を昇順に格納した配列に対して、行きがけ順の配列の要素uの値(クエリ)を用いてlower_bounds
   

返りかけ順の復習になった

#include <iostream>
#include <string>
#include <map>
#include <unordered_map>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <vector>
#include <queue>
#include <stack>
#include <limits.h>
#include <bitset>
#include <list>
#include <set>
#include <numeric>
#include <tuple>

int N;
int Q;
std::vector<int> P;
std::vector<std::vector<int>> G;
std::vector<int> go;
std::vector<int> back;
std::vector<std::vector<int>> depth; // 深さ毎に行きがけ順番号を昇順に格納していく
std::vector<bool> seen;

int count = 0;

void dfs(const int& u, int d) {
	// 訪問
	seen[u] = true;

	// 行きがけ順
	go[u] = count;
	count++;

	// 深さ管理
	depth[d].push_back(go[u]);

	// 隣接点
	for (const int& v : G[u]) {
		if (seen[v] == true) {
			continue;
		}
		
		dfs(v, d + 1);
	}

	// 帰りがけ順
	back[u] = count;
	count++;
}

int main()
{
	std::cin.tie(0);
	std::ios::sync_with_stdio(false);

	std::cin >> N;

	G.resize(N);
	go.resize(N);
	back.resize(N);
	depth.resize(N);
	seen.resize(N);

	for (int i = 2; i <= N; i++) {
		int p;
		std::cin >> p;
		p--;
		P.push_back(p);

		// 有効グラフ
		int c = i - 1;

		// 親から子へ
		G[p].push_back(c);
	}

	// dfsで行きがけ順、帰りがけ順の記録
	int s = 0;
	int d = 0;
	dfs(s, d);

	std::cin >> Q;

	std::vector<int> vecAns;
	for (int i = 0; i < Q; i++) {
		int u, d;
		std::cin >> u >> d;
		u--;

		// 深さ毎に行きがけ順番号を昇順に格納した配列に対して、帰りがけ順の配列の要素uの値を用いてlower_bounds
		// 深さ毎に行きがけ順番号を昇順に格納した配列に対して、行きがけ順の配列の要素uの値を用いてlower_bounds
		auto ans = std::lower_bound(depth[d].begin(), depth[d].end(), back[u]) - std::lower_bound(depth[d].begin(), depth[d].end(), go[u]);
		vecAns.push_back(ans);
	}

	for (int i = 0; i < vecAns.size(); i++) {
		std::cout << vecAns[i] << std::endl;
	}
	
	return 0;
}