Atcoder解いてみた AtCoder Beginner Contest D - 閉路

AtCoder Beginner Contest D - 閉路

問題はこちら

回答

クエリによって結ばれる2点のLCAを求め、2点とLCAによって構成される三角形を構成する頂点数をカウントする
しかし、素直にやると(QN)でTLE

解説では、以下のように解いていた
DFSで木を構成する各点の親、深さを求める
ダブリングを用いて、LCAを求める
ダブリング表を作った後のLCA算出の際、何回ダブリング表を用いて木をたどればいいかわからないため、LCAを越さないギリギリのところまで移動するための実装の工夫がとても勉強になる

計算量だが、
DFSにO(N)
ダブリング表作成にO(N logN)
クエリの度にダブリング表を持ちてたどるのにO(Q logN)
なため、計算量はO((Q+N)log(N))

#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, Q;

std::vector<std::vector<int>> G;
std::vector<int> depth;
std::vector<int> parent;

void dfs(const int& u, int d) {
	depth[u] = d;
	for (const int & v : G[u]) {
		if (depth[v] != -1) {
			continue;
		}

		parent[v] = u;
		dfs(v, d+1);
	}
}

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

	std::cin >> N;
	
	// 頂点1を根として、木を構築する
	G.resize(N);
	depth.resize(N, -1);
	parent.resize(N);

	for (int i = 0; i < N-1; i++) {
		int x, y;
		std::cin >> x >> y;
		x--;
		y--;

		G[x].push_back(y);
		G[y].push_back(x);
	}

	// DFSで最短距離(深さ)と最短距離を求める
	int s = 0;
	int d = 0;
	dfs(s, d);

	// ダブリング
	// parentが1回移動した先

	// 移動回数K <= 10^6程度まで対応
	const int log = 20;

	// doubling[k][v] : 頂点vから2^k回移動した先の頂点
	std::vector<std::vector<int>> doubling(log, std::vector<int>(N));

	for (int v = 0; v < N; v++) {
		// 1回移動した先
		doubling[0][v] = parent[v];
	}


	// 2^k回移動
	// = 2^(k-1)回移動後、さらに2^(k-1)回移動
	for (int k = 1; k < log; k++) {
		for (int v = 0; v < N; v++) {
			doubling[k][v] = doubling[k - 1][doubling[k - 1][v]];
		}
	}

	std::cin >> Q;

	std::vector<int> vecAns;
	for (int q = 0; q < Q; q++) {
		int a, b;
		std::cin >> a >> b;
		a--;
		b--;

		int oa = a;
		int ob = b;

		// a,bの深さをそろえる
		int* adj = nullptr;
		int adjNum = std::abs(depth[a] - depth[b]);
		if (depth[a] > depth[b]) {
			adj = &a;
		}
		if (depth[a] < depth[b]) {
			adj = &b;
		}

		// 高さをそろえる
		if (adj != nullptr) {
			for (int k = 0; k < log; k++) {
				if ((adjNum >> k) & 1LL) {
					*adj = doubling[k][*adj];
				}
			}
		}

		// a,bからそれぞれ根までさかのぼり、LCAを求める
		int lca;
		if (a == b) {
			lca = a;
		}
		else {

			// 以下の実装だと、lca飛び越えて、根まで行ってしまう可能性がある
			// →lcaの直下まで飛ぶように工夫する
			//for (int k = 0; k < log; k++) {
			//	if ((N >> k) & 1LL) {
			//		a = doubling[k][a];
			//		b = doubling[k][b];

			//		if (a == b) {
			//			lca = a;
			//			break;
			//		}
			//	}
			//}

			// 大きいジャンプから試す
			// 以下のコードにより、大きく飛びすぎてLCAより上の同じノードにたどりついても更新されない
			// a,bがLCA以下の限り位置が更新され、最終的にa,bはLCA直下に至る
			for (int k = log - 1; k >= 0; k--) {

				// 2^k上に行ってもまだ別なら上がる
				if (doubling[k][a] != doubling[k][b]) {
					a = doubling[k][a];
					b = doubling[k][b];
				}
			}

			// この時点でa,bはLCAの直下
			lca = doubling[0][a];
		}

		int ans = depth[oa] + depth[ob] - 2 * depth[lca] + 1;
		vecAns.push_back(ans);
	}

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