mayoko’s diary

プロコンとかいろいろ。

Educational Codeforces Round 7 F. The Sum of the k-th Powers

解法

k 乗の和は, k+1 乗の多項式になります。なんで, って人は下の記事を読みましょう。
yama-taku.science

ということで, ラグランジュ補間しましょう。
mayokoex.hatenablog.com
ラグランジュ補間は次数 k に対して O(k^2) かかるのが普通ですが, 今回の場合は 0, 1, 2, ... と規則正しく既知の値を取ってこれるので,

 Q_i(i) = i * Q_{i-1}(i-1) / (i-1-k), Q_i(n) = Q_{i-1}(n) * (n-i+1) / (n-i)

と O(1) で更新できます。これは過去の AtCoder でも出題されてますね。

vector<int> G[500050];

void dfs(int v, int p, int d, vi& dist) {
    bool leaf = true;
    for (int ch : G[v]) {
        if (ch==p) continue;
        leaf = false;
        dfs(ch, v, d+1, dist);
    }
    if (leaf) dist.push_back(d);
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int n;
    cin >> n;
    for (int i = 0; i < n-1; i++) {
        int x, y;
        cin >> x >> y;
        x--; y--;
        G[x].push_back(y);
        G[y].push_back(x);
    }
    int ans = 0;
    for (int v : G[0]) {
        vector<int> dist;
        dfs(v, 0, 0, dist);
        sort(dist.begin(), dist.end());
        int size = dist.size();
        for (int i = 0; i < size-1; i++) {
            dist[i+1] = max(dist[i]+1, dist[i+1]);
        }
        ans = max(dist[size-1]+1, ans);
    }
    cout << ans << endl;
    return 0;
}