けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

Codeforces Round #614 (Div. 1) C. Xenon's Attack on the Gangs (R2300)

面白かった!!!こういうのを確実に通せるようにならないと!!!

問題へのリンク

問題概要

 N 頂点の木が与えられる。木の各辺を  0, 1, \dots, N-2 のラベルをつける方法のうち、

  •  \sum_{1 \le u \lt v \le N} f(u, v)

の値の最大値を求めよ。ただし  f(u, v) は、2 頂点  u, v を結ぶパスに含まれる辺の値の集合を考えたときに、それに含まれない最小の 0 以上の整数値を表す。

制約

  •  2 \le N \le 3000

考えたこと

 f(u, v) の意味が釈然としないので、求める値を、次のように言い換えてみよう。

  • 0 の辺を含むパスの本数
  • 0 の辺と 1 の辺をともに含むようなパスの本数
  • 0 の辺と 1 の辺と 2 の辺をすべて含むようなパスの本数
  • ...

を合計した値に等しくなる。ここまでは少し考えればわかる。

まず 0 を含むパスの本数というのは、「0 の辺」でツリー全体を 2 つに分離したときに、それぞれの部分木のサイズの積に等しい。

同様にして、0 の辺と 1 の辺をともに含むようなパスの本数は、それらの辺をともに含むようなパスをとってきたときに、その両端点を根とする部分木のサイズの積に等しい。

観察

上図では、0 の辺と 1 の辺を引き離して描いたけど、これは明らかに損である。数値を入れ替えて、0 と 1 が隣接するようにした方が得になる。

同様にして、0, 1, 2, ... とラベルを増やしていくときに、新たな値をつける辺を見定めるときには、「すでに値を振ったパス」の両端に接続するようにするのが良いということがいえる。ただし葉に到達したら終了。それ以降は何をやっても値は増えない。

ちゃんとした証明は editorial に。

DP へ

以上のことから、次のような DP が立つ

  • dp[ u ][ v ] := 無の状態から開始して、パス (u, v) まで伸ばす過程で合算されたスコアの最大値

これは、「パスを伸ばす直前がどうだったか」で、以下の 2 通りの場合分けをして遷移できる。ここで s(u, v) を、v を根とし、u を親方向の頂点とした場合の v を根とする部分木のサイズとする。また、

  • u-v パスにおいて、u から v の方向へ 1 辺進んだ点を up
  • u-v パスにおいて、v から u の方向へ 1 辺進んだ点を vp

とする。このとき、遷移は以下のように表せる。

u-v パスが、up-v パスから辺 (up, u) を伸ばすことで成長した場合

chmax(dp[ u ][ v ], dp[ up ][ v ] + s(up, u) × s(vp, v))

u-v パスが、u-vp パスから辺 (vp, v) を伸ばすことで成長した場合

chmax(dp[ u ][ v ], dp[ u ][ vp ] + s(up, u) × s(vp, v))

計算量

まとめると、以下の値を前処理して求めておけば、 O(N^{2}) な DP で求めることができる。

  • s(u, v)
  • u-v パスにおいて、u から v 方向に 1 辺進んだときの頂点

これらはライブラリにしてしまうことにした。前処理は  O(N\log{N}) を要する。上記のうちの後者は LCA を求める要領で  O(\log{N}) で求める。この場合、全体の計算量は  O(N^{2}\log{N}) となる。

ただし今回は前処理に  O(N^{2}) かけてよく、上記の値をすべてメモリに置いておくこともできるので、全体として  O(N^{2}) で解くことができる (公式解答)。

#include <iostream>
#include <vector>
#include <map>
using namespace std;

using Graph = vector<vector<int> >;
struct RunTree {
    // id[v][w] := the index of node w in G[v]
    vector<map<int,int> > id;

    // num[v][i] := the size of subtree of G[v][i] with parent v
    vector<vector<long long> > num;

    // size(u, v) := the size of subtree v with parent u
    long long size(int u, int v) {
        return num[u][id[u][v]];
    }

    // lca(u, v)
    int getLCA(int u, int v) {
        if (depth[u] > depth[v]) swap(u, v);
        for (int i = 0; i < (int)parent.size(); ++i)
            if ( (depth[v] - depth[u]) & (1<<i) )
                v = parent[i][v];
        if (u == v) return u;
        for (int i = (int)parent.size()-1; i >= 0; --i) {
            if (parent[i][u] != parent[i][v]) {
                u = parent[i][u];
                v = parent[i][v];
            }
        }
        return parent[0][u];
    }

    // length(u, v)
    long long length(int u, int v) {
        int lca = getLCA(u, v);
        return depth[u] + depth[v] - depth[lca]*2;
    }

    // getParent(v, p) := the parent of v directed for p
    int getParent(int v, int p) {
        if (v == p) return -1;
        int lca = getLCA(v, p);
        if (lca != v) return parent[0][v];
        for (int i = (int)parent.size()-1; i >= 0; --i) {
            if (parent[i][p] != -1 && depth[parent[i][p]] > depth[v]) {
                p = parent[i][p];
            }
        }
        return p;
    }
    
    // rec
    vector<vector<int> > parent;
    vector<int> depth;
    int rec(const Graph &G, int v, int p = -1, int d = 0) {
        int p_index = -1;
        int sum = 1;
        parent[0][v] = p;
        depth[v] = d;
        for (int i = 0; i < (int)G[v].size(); ++i) {
            int ch = G[v][i];
            id[v][ch] = i;
            if (ch == p) {
                p_index = i;
                continue;
            }
            int s = rec(G, ch, v, d+1);
            num[v][i] = s;
            sum += s;
        }
        if (p_index != -1) num[v][p_index] = (int)G.size() - sum;
        return sum;
    }

    // init
    void init(const Graph &G) {
        int N = (int)G.size();
        id.assign(N, map<int,int>());
        num.assign(N, vector<long long>());
        for (int v = 0; v < N; ++v) num[v].assign((int)G[v].size(), 0);
         int V = (int)G.size();
        int h = 1;
        while ((1<<h) < N) ++h;
        parent.assign(h, vector<int>(N, -1));
        depth.assign(N, -1);
        rec(G, 0);
        for (int i = 0; i+1 < (int)parent.size(); ++i)
            for (int v = 0; v < V; ++v)
                if (parent[i][v] != -1)
                    parent[i+1][v] = parent[i][parent[i][v]];
    }
};



const long long INF = 1LL<<60;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

using pint = pair<int,int>;
int N;
Graph G;
RunTree rt;

vector<vector<long long> > dp;
long long dprec(int u, int v) {
    if (dp[u][v] != INF) return dp[u][v];
    if (dp[v][u] != INF) return dp[v][u];
    if (u == v) return 0;

    long long res = 0;
    int up = rt.getParent(u, v);
    int vp = rt.getParent(v, u);
    chmax(res, dprec(up, v));
    chmax(res, dprec(vp, u));
    res += rt.size(up, u) * rt.size(vp, v);
    return dp[u][v] = dp[v][u] = res;
}

long long solve() {
    rt.init(G);
    dp.assign(N, vector<long long>(N, INF));
    long long res = 0;
    for (int i = 0; i < N; ++i) {
        for (int j = i+1; j < N; ++j) {
            chmax(res, dprec(i, j));
        }
    }
    return res;
}

int main() {
    while (scanf("%d", &N) != EOF) {
        G.assign(N, vector<int>());
        for (int i = 0; i < N-1; ++i) {
            int u, v;
            scanf("%d %d", &u, &v);
            --u, --v;
            G[u].push_back(v);
            G[v].push_back(u);
        }
        cout << solve() << endl;
    }
}