はじめに

この記事はtraPアドベントカレンダー2020の30日目(12/13)の記事です。
19Bのebiです。普段はAtCoderで競技プログラミングを楽しんでいます(ID:ebi_fly)。
クリスマスといえばクリスマスツリー！ということでここでは全方位木DPのアルゴリズムを解説し、具体的に全方位木DPを用いて木の直径を求めてみます。

全方位木DPってなに？

簡単に言うと条件を満たす一つの頂点を根とした根付き木に対して行った木DPを全頂点について求めるアルゴリズムです。
「木」や「DFS」などの単語になじみのない人はこの記事読んでみてください。

全方位木DPのアルゴリズム

全方位木DPは以下のようにして計算します。

1. ある一つの頂点を根とした根付き木に対して木DPを行う。
2. 1.の結果からDFS順に全ての頂点に対してその頂点を根としたときの木DPの結果を計算する。

ここで2.の箇所を効率的に計算することで愚直では$O(N^2)$かかるところを$O(N)$で動作させます。
以下では木の直径を求めることを通して全方位木DPについて見ていきます。
木の直径を求めるにはある一つの頂点を根としたときの根付き木に対してその頂点から最も遠い頂点との距離を求めます。次に、根とした頂点を経由して最も距離が遠い2頂点間の距離を求め、これを全頂点に対して行いそれらの最大値が木の直径となります。

木DPについて

ステップ1.をおこないます。木DPでは多くの場合現在いる頂点の子を根とした部分木について木DPを再帰的におこないます。これは根からDFSを帰りがけで処理することでできます。
今回は部分木に対して最も遠い頂点の距離を木DPで求めます。これは以下のようなdpを根からDFSを帰りがけで処理することでできます。$v$は現在いる頂点、$child$は$v$の子です。

$$dp[v]=\underset{child\in tree[v]}{\max }(dp[child]+cos{t}_{v,child})$$

以下が実装例です。後で部分木についての結果を用いるので配列dpに保存しています。

void dfs1(int v, int par){
    dp[v] = 0;
    for(auto [nv, cost]: tree[v]){
        if(nv==par) continue;
        dfs1(nv,v);
        dp[v] = std::max(dp[v],dp[nv]+cost);
    }
    return;
}

全頂点に対してDFSをおこなう

ステップ2.をおこないます。普通に全頂点に対してステップ1.と同様の木DPを行いたいですがそれだと計算量が$O(N^2)$になってしまいます。そこでステップ1.の処理で作られた部分木についての結果を用いて高速化します。
具体的に考えてみます。頂点$v$についてみると子ノードについての木DPはステップ1.で完了しているので親からの結果さえあればあとはそれをマージすることで$v$を根としたときの木DPの結果を求めることができます。仮に親からの結果を$val$とすると木DPの結果$result$は以下のようになります。ここで$op$は子の結果をマージする関数です。

$$result=op(dp[v],val)$$

$v$から最も遠い頂点との距離を求めたいときは以下のようになります。

$$result=max(dp[v],val+cos{t}_{v,parent})$$

あとは親からの結果さえわかればよいです。$v$の親を根とし$v$の部分だけ取り除いた木についての木DPの結果が$val$になっています。
DFS順に処理をすることで親ノードの処理が終わった後にそのノードの処理をすることができるためDFS順に処理をします。また、次の頂点に遷移する際に遷移先の頂点を取り除いた木DPの結果を与えることで上に記したような処理ができます。
以下が実装例です。$v$が探索する頂点、$par$が親、$val$が$v$から$par$に木DPしたときの結果です。

void dfs2(int v, int par, ll val){
    std::vector<std::pair<ll,int>> ch;
    ch.emplace_back(std::make_pair(0,-1));
    //vの子を見てvと子の部分木に含まれる最も遠い頂点までの距離をchに格納していく
    for(auto [nv, cost]: tree[v]){
        if(nv==par){
            ch.emplace_back(std::make_pair(val+cost,par));
            continue;
        }
        ch.emplace_back(std::make_pair(dp[nv]+cost,nv));
    }
    std::sort(ch.rbegin(),ch.rend());
    dtr = std::max(dtr,ch[0].first+ch[1].first);//直径を求める
    for(auto e: tree[v]){
        int nv = e.to;
        if(nv==par) continue;
        if(nv==ch[0].second){
            dfs2(nv,v,ch[1].first);
            continue;
        }
        dfs2(nv,v,ch[0].first);
    }
    return;
}

完成！

これで木の直径を求めることができました。提出
二つのステップを考えることで他の全方位木DPができると思います。
また全方位木DPには抽象化したライブラリも存在するので興味がある人はこの記事を読んでみるとよいと思います。

終わりに

全方位木DPは理解するのがむずかしいと思いますがとてもおもしろいアルゴリズムなのでぜひトライしてみてください。
明日は@Hinaruhi さんの記事です。お楽しみに！