Cod sursa(job #2249504)

Utilizator TincaMateiTinca Matei TincaMatei Data 29 septembrie 2018 22:53:29
Problema Lowest Common Ancestor Scor 70
Compilator cpp Status done
Runda Arhiva educationala Marime 3.8 kb
#include <bits/stdc++.h>

using namespace std;

struct Tree {
	int N;
	std::vector<std::vector<int> > G;
	
	// Basic tree statistics
	int *label;
	int *subarb, *depth;
	int lastStatisticRoot;

	Tree(int _N) : N(_N) {
		G.resize(1+N);
		label = subarb = depth = NULL;
		lastStatisticRoot = 0;
	}

	void pushEdge(int a, int b) {
		G[a].push_back(b);
		G[b].push_back(a);
	}
	
	void statisticDfs(int root = 1, int father = 0) {
		subarb[root] = 1;
		for(auto it: G[root])
			if(it != father) {
				depth[it] = depth[root] + 1;
				statisticDfs(it, root);
				subarb[root] += subarb[it];
			}
	}

	void calculateStatistics(int root = 1) {
		if(subarb == NULL && depth == NULL) {
			subarb = new int[1+N];
			depth = new int[1+N];
		}

		if(lastStatisticRoot != root) {
			memset(subarb, 0, sizeof(int) * (1 + N));
			memset(depth,  0, sizeof(int) * (1 + N));
	
			statisticDfs(root);
		
			lastStatisticRoot = root;
		}
	}

	void allocLabels() {
		if(label == NULL)
			label = new int[1+N];
		memset(label, 0, sizeof(int) * (1 + N));
	}
};

struct BinaryLifting {
	Tree *T;
	int N;
	int maxlg, lastRoot;
	int **ancestor;
	
	BinaryLifting(int n) {
		N = n;
		maxlg = 0;
		while((1 << maxlg) < N)
			++maxlg;
		ancestor = new int*[1+maxlg];
		for(int i = 0; i <= maxlg; ++i) {
			ancestor[i] = new int[1 + N];
			memset(ancestor[i], 0, sizeof(int) * (1 + N));
		}
	}

	BinaryLifting(Tree *_T):T(_T) {
		N = T->N;
		maxlg = 0;
		while((1 << maxlg) < T->N)
			++maxlg;
		
		ancestor = new int*[1+maxlg];
		for(int i = 0; i <= maxlg; ++i) {
			ancestor[i] = new int[1 + T->N];
			memset(ancestor[i], 0, sizeof(int) * (1 + T->N));
		}
		
		lastRoot = 0;
	}

	void ancestryDfs(int nod, int papa = 0) {
		ancestor[0][nod] = papa;
		for(int i = 1; i <= maxlg; ++i)
			ancestor[i][nod] = ancestor[i - 1][ancestor[i - 1][nod]];
		for(auto it: T->G[nod])
			if(it != papa)
				ancestryDfs(it, nod);
	}

	void setAncestor(int nod, int anc) {
		ancestor[0][nod] = anc;
	}

	void resetAncestor(int nod, int anc) {
		ancestor[0][nod] = anc;
		for(int i = 1; i <= maxlg; ++i)
			ancestor[i][nod] = ancestor[i - 1][ancestor[i - 1][nod]];
	}

	void ancestryLiniar() {
		for(int i = 1; i <= maxlg; ++i)
			for(int nod = 1; nod <= N; ++nod)
				ancestor[i][nod] = ancestor[i - 1][ancestor[i - 1][nod]];
	}

	void computeAncestry(int root = 1) {
		if(T == NULL) {
			ancestryLiniar();
		} else if(lastRoot != root) {
			lastRoot = root;
			ancestor[0][root] = 0;
			ancestryDfs(1);
		}
	}

	int goUp(int nod, int x) {
		if(x >= N)
			return 0;
		for(int i = 0; i <= maxlg; ++i)
			if((1 << i) & x)
				nod = ancestor[i][nod];
		return nod;
	}
};

struct LCA {
	BinaryLifting *bl;
	Tree *T;
	int lastRoot;

	LCA(Tree *_T):T(_T) {
		lastRoot = 0;
		bl = new BinaryLifting(T);
	}

	void computeAncestry(int root) {
		if(root != lastRoot) {
			bl->computeAncestry(root);
			T->calculateStatistics(root);
			lastRoot = root;
		}
	}
	
	int getLca(int a, int b) {
		if(T->depth[a] > T->depth[b])
			a = bl->goUp(a, T->depth[a] - T->depth[b]);
		else
			b = bl->goUp(b, T->depth[b] - T->depth[a]);

		for(int i = bl->maxlg; i >= 0; --i)
			if(bl->ancestor[i][a] != bl->ancestor[i][b]) {
				a = bl->ancestor[i][a];
				b = bl->ancestor[i][b];
			}
		if(a != b)
			a = bl->ancestor[0][a];
		return a;
	}
};

int main() {
#ifdef HOME
	FILE *fin = fopen("input.in", "r");
	FILE *fout = fopen("output.out", "w");
#else
	FILE *fin = fopen("lca.in", "r");
	FILE *fout = fopen("lca.out", "w");
#endif
	
	int n, q;
	Tree *T;
	LCA *lcaSolver;

	fscanf(fin, "%d%d", &n, &q);
	T = new Tree(n);
	
	for(int i = 2; i <= n; ++i) {
		int x;
		fscanf(fin, "%d", &x);
		T->pushEdge(i, x);
	}
	lcaSolver = new LCA(T);
	lcaSolver->computeAncestry(1);

	for(int i = 0; i < q; ++i) {
		int a, b;
		fscanf(fin, "%d%d", &a, &b);
		fprintf(fout, "%d\n", lcaSolver->getLca(a, b));
	}

#ifdef HOME
	fclose(fin);
	fclose(fout);
#endif
	return 0;
}