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

問題

indeednow-finala-open.contest.atcoder.jp

解法

木DP。

コード

#include <cstdio>
#include <iostream>
#include <sstream>
#include <fstream>
#include <iterator>
#include <iomanip>
#include <algorithm>
#include <cmath>
#include <string>
#include <vector>
#include <list>
#include <queue>
#include <stack>
#include <set>
#include <map>
#include <bitset>
#include <numeric>
#include <limits>
#include <climits>
#include <cfloat>
#include <functional>
using namespace std;

typedef long long ll;

struct Edge {
  int to, type;
  Edge(int to, int type) {
    this->to = to;
    this->type = type;
  }
};

vector<vector<Edge>> tree;
vector<vector<int>> dp;

int dfs(int v, int p) {
  int res = 0;
  for (Edge e : tree[v]) {
    if (e.to == p) continue;
    res = max(res, dfs(e.to, v));

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

  return res;
}

int main() {
  int N;
  cin >> N;
  tree.resize(N);
  for (int i = 0; i < N - 1; ++i) {
    int a, b, t;
    cin >> a >> b >> t;
    a--, b--;
    if (t == 1) {
      tree[a].push_back(Edge(b, 0));
      tree[b].push_back(Edge(a, 1));
    } else {
      tree[a].push_back(Edge(b, 2));
      tree[b].push_back(Edge(a, 2));
    }
  }
  dp.assign(2, vector<int>(N));
  cout << dfs(0, -1) << endl;
}