NoiminのNoise

競技プログラミング (多め) とWeb (たまに) ,自然言語処理 (ブログではまだ)。数式の書き方を一気に KaTeX に変えようとして記事を全削除してインポートし直すなどしたので,過去にブックマークされた記事は URL が変わってしまっている可能性があります…….

Codeforces Round #588 (Div. 1 B / Div 2. E) Kamil and Making a Stream

問題: https://codeforces.com/contest/1229/problem/B

公式解説: https://codeforces.com/blog/entry/70008

問題概要

$ n (\le {10}^5) $ 頂点の木が与えられる.$ i $ 番目の頂点には非負整数 $ a_i (\le {10}^{12} )$ が割り当てられている. 次式の値を求めよ.

$ \sum_{u は v の先祖} f(u, v) $

ただし $ f(u, v) $ は $ u $ から $ v $ までのパス上にあるすべての頂点に割り当てられている非負整数の最大公約数を指す.

解法概要

一言で言ってしまえば,「上から順に gcd の値を保持しながら dfs していく」それだけ.

実はこの解法,コンテスト中に思いついたものの秒で棄却してしまった解法だった.

「木のパスに分岐がないような最悪ケースの場合,空間計算量 $ O(n) $ の gcd を保持してノードごとに更新していく必要があり,時間計算量が $ O(n^2) $ になってしまうのでは?」

と考えたためである.

しかしこれは誤りで,gcd の種類数は $ \log (\max a_i) $ で抑えることができる.何故ならば,根からパスをたどっていく過程で gcd の値が変化するとき,新しい gcd $ g' $ と古い gcd $ g $ について, $ g' \le \frac{g}{2} $ が成り立つためである.$ g' = \mathit{gcd}(g, a_i )$ なので, $ g' $ は必ず $ g $ の約数でもある.なので,$ g' \le g $ が成り立つ場合,$ g' \le \frac{g}{2} $ も成り立つことになる.このことから,gcd の種類数は $ \log (\max a_i) $ で抑えることができることがわかる.

ソースコード

#include <iostream>
#include <vector>
#include <map>

using namespace std;

using ll =  long long;

constexpr ll MOD = 1000000007;
constexpr int N_MAX = 100000;

int n, u, v;
vector<int> graph[N_MAX];
vector<ll> nodes(N_MAX);

ll gcd(ll a, ll b){
    if(b == 0) return a;
    else return gcd(b, a%b);
}

ll dfs(int node, int par, map<ll, ll> mp) {
    ll ret = 0;
    map<ll, ll> mp2;
    for(auto p: mp) {
        ll g = gcd(nodes[node], p.first);
        mp2[g] += p.second;
        ret += (g * p.second) % MOD;
        ret %= MOD;
    }
    ++mp2[nodes[node]];
    ret += nodes[node];
    ret %= MOD;
    for(int child: graph[node]) {
        if(child == par) continue;
        ret += dfs(child, node, mp2);
        ret %= MOD;
    }
    return ret;
}


int main() {
    std::ios::sync_with_stdio(false); cin.tie(nullptr);
    cin >> n;
    for(int i=0;i<n;++i) cin >> nodes[i];
    for(int i=0;i<n-1;++i) {
        cin >> u >> v;
        --u; --v;
        graph[u].push_back(v);
        graph[v].push_back(u);
    }

    cout << dfs(0, -1, map<ll, ll>()) << endl;
}

所感

計算量解析,大事。