Indeed なう A日程 D: Longest Path (木DP)

問題

D: Longest Path - Indeedなう（オープンコンテスト） | AtCoder

解法

ある頂点vに注目した時に、「vに入ってくる最大のパス」と「vから出て行く最大のパス」を計算する。この2つが同じパスにならないように、頂点はそれぞれ1度ずつしか注目しない。

めっちゃ理解しづらかった。

コード

import java.io.IOException;
import java.util.ArrayList;

class Main {
	private static boolean used[];
	private static ArrayList<ArrayList<Edge>> list;
	private static int[][] dp;
	private static final int OUT = 0;
	private static final int IN = 1;

	public static void main(String[] args) {
		int N = nextInt();
		list = new ArrayList<>();
		for (int i = 0; i < N; i++) {
			list.add(new ArrayList<Edge>());
		}

		for (int i = 0; i < N - 1; i++) {
			int a = nextInt() - 1;
			int b = nextInt() - 1;
			int t = nextInt();

			if (t == 1) {
				list.get(a).add(new Edge(b, 0));// a->b
				list.get(b).add(new Edge(a, 1));// a->b
			} else {
				list.get(a).add(new Edge(b, 2));
				list.get(b).add(new Edge(a, 2));
			}

		}

		used = new boolean[N];

		// dp[0][v] := 頂点vから出る最長パス
		// dp[1][v] := 頂点vに入る最長パス
		dp = new int[2][N];

		System.out.println(dfs(0));

	}

	static int dfs(int from) {
		used[from] = true;
		int res = 0;

		for (Edge e : list.get(from)) {
			int to = e.to;

			if (used[to]) {
				continue;
			}
			res = Math.max(res, dfs(to));

			if (e.type == 0 || e.type == 2) {
				// v -> e.to
				res = Math.max(res, dp[IN][from] + dp[OUT][to]);
			}
			if (e.type == 1 || e.type == 2) {
				res = Math.max(res, dp[OUT][from] + dp[IN][to]);
			}
			if (e.type == 0 || e.type == 2) {
				// v -> e.to
				dp[OUT][from] = Math.max(dp[OUT][from], dp[OUT][to] + 1);
			}
			if (e.type == 1 || e.type == 2) {
				dp[IN][from] = Math.max(dp[IN][from], dp[IN][to] + 1);
			}

		}
		return res;

	}

	static class Edge {
		int to, type;

		public Edge(int to, int type) {
			this.to = to;
			this.type = type;
		}
	}

	static int nextInt() {
		int c;
		try {
			c = System.in.read();
			while (c != '-' && (c < '0' || c > '9'))
				c = System.in.read();
			if (c == '-')
				return -nextInt();
			int res = 0;
			while (c >= '0' && c <= '9') {
				res = res * 10 + c - '0';
				c = System.in.read();
			}
			return res;
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		return -1;
	}
}