Atcoder解いてみた AtCoder Beginner Contest C - Ameba

AtCoder Beginner Contest C - Path Graph?

問題はこちら

回答

以下のコードではTLE
→木が一直線になると、O(N^2)

#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>
#include <iomanip>
#include <random>

int N;

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

	std::cin >> N;

	// parent[k]: アメーバ k の親の番号
	// アメーバの最大番号は 2*N + 1
	std::vector<int> parent(2 * N + 2, 0);

	for (int i = 1; i <= N; i++) {
		int a;
		std::cin >> a;

		parent[2 * i] = a;
		parent[2 * i + 1] = a;
	}

	for (int i = 1; i <= 2 * N + 1; i++) {
		int c = i;
		int ans = 0;
		
		// アメーバ1までたどる
		while (c != 1) {
			c = parent[c];
			ans++;
		}

		std::cout << ans << std::endl;
	}


	return 0;
}

以下のように、親の番号が入力されたときに、子の親の数をカウントする

#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>
#include <iomanip>
#include <random>

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

	std::cin >> N;

	// depth[k]: アメーバ k の上にどれだけ親がいるか
	// アメーバの最大番号は 2*N + 1
	std::vector<int> depth(2 * N + 2, -1);

	depth[1] = 0;
	for (int i = 1; i <= N; i++) {
		int a;
		std::cin >> a;

		depth[2 * i] = depth[a]+1;
		depth[2 * i + 1] = depth[a] + 1;
	}

	for (int i = 1; i <= 2 * N + 1; i++) {
		std::cout << depth[i] << std::endl;
	}


	return 0;
}