(木DP・数え上げ)ABC-036-D-塗り絵

問題URL→http://abc036.contest.atcoder.jp/tasks/abc036_d

木DPの典型問題。
int rec(int v, int p, int color)
とすると計算量がオーバーするのでpを消去するためにdfsを一回走らせて前の頂点に戻る辺を除く。

数え上げでコード上で*になっていたところが+でやったら答えが合わなかった。
数え上げなので掛けなければならない。ところどころ+のところもあるのでよく考えよう。

木DP何気に初めて実装したけどうまくかけたと思う。他の人のコードも見てみよう。

#include <iostream>
#include <vector>
#include <string.h>
#include <stack>
#include <queue>
#include <algorithm>
#include <climits>
#include <cmath>
#include <map>
#define REP(i,n) for(ll i=0;i<(n);i++)
#define MOD 1000000007
#define int long long
#ifdef int
const long long INF = LLONG_MAX / 10;
#else
const int INF = 1010101010;
#endif
using namespace std;
typedef long long ll;
typedef vector<vector<ll>> mat;
typedef pair<int, int> P;
//typedef pair<double, double> P;

int N;
int a, b;

vector<int> G[101010];
vector<int> GG[101010];

void dfs(int v, int p)
{
    for (int ne : G[v]) {
        if (ne == p) continue;
        GG[v].push_back(ne);
        dfs(ne, v);
    }
}

int dp[101010][2];

int rec(int v, int color)
{
    if (dp[v][color] != -1) return dp[v][color];

    int res = 1;
    for (int ne : GG[v]) {
        if (color == 0) {
            res *= rec(ne, 1);
        } else {
            res *= rec(ne, 0) + rec(ne, 1);
        }
        res %= MOD;
    }
    return dp[v][color] = res;
}

signed main()
{
    cin >> N;
    REP(i,N-1) {
        cin >> a >> b;
        a--; b--;
        G[a].push_back(b);
        G[b].push_back(a);
    }

    dfs(0, -1);

    memset(dp, -1, sizeof dp);
    cout << (rec(0, 0) + rec(0, 1)) % MOD << endl;
}

(木の直径・木DP)ARC-022-C-ロミオとジュリエット

問題URL→http://arc022.contest.atcoder.jp/tasks/arc022_3

知らない典型問題っぽかったのですぐ解説見た。
公式の解説の木DPが良く分からなかったので次のkmjpさんのブログを参考にした。
DFSを2回行うのが簡単らしいが木DPのほうが応用が利くらしい。

#include <bits/stdc++.h>
#define FOR(i,bg,ed) for(ll i=(bg);i<(ed);i++)
#define REP(i,n) FOR(i,0,n)
#define MOD 1000000007
#define int long long
using namespace std;
typedef long long ll;
const int INF = 1e9;

int N;
int x, y;
vector<int> E[101010];
int sx, sy;
int ma;

pair<int, int> dfs(int cur, int pre)
{
    vector<pair<int, int>> V;

    REP(i,E[cur].size()) {
        int tar = E[cur][i];
        if (tar == pre) continue;
        V.push_back(dfs(tar, cur));
    }

    V.push_back(make_pair(1, cur));

    sort(V.begin(), V.end());
    reverse(V.begin(), V.end());
    if (V.size() >= 2) {
        int tar = V[0].first + V[1].first;
        if (tar > ma) {
            ma = tar;
            sx = V[0].second + 1;
            sy = V[1].second + 1;
        }
    }

    return make_pair(V[0].first + 1, V[0].second);
}

signed main()
{
    cin >> N;
    REP(i,N-1) {
        cin >> x >> y;
        x--; y--;
        E[x].push_back(y);
        E[y].push_back(x);
    }

    dfs(0, -1);
    cout << sx << " " << sy << endl;
}