ふみふみのプログラミング自習室

私(プログラム初心者)が最近勉強したことをまとめています。

ABC248-Gの包除原理を用いた解き方について

書くためのネタがあまりなかったのですが、更新できてない方が問題なので、とりあえず書きました。

ABC248-Gの包除原理を用いた解答

 D = 10 ^ 5以下の自然数における約数の個数の最大値」としたとき、

包除原理による時間計算量/空間計算量 O(DN) の解答例を紹介する。


 total weights[v ]  = 相違なる 2 点間のパスの gcd の値が vとなるようなパスに含まれる頂点数の総和

と定義する。

すると、答えは以下の式で表せることがわかる。

 \displaystyle

ans =\sum_{v=1}^{10^5} v  ・ total weights[v]

ここで包除原理を用いると、 S_i を「異なる 2 点間のパスの gcdの値が i の倍数となるようなパスに含まれる頂点数の総和」とし

 \displaystyle

total weights[i] = S_i - \sum_{i | j ,  i < j}^{}total weights[j]

である。

これより、任意の  i(1  \leq  i  \leq 10 ^ 5)に対して上記で定めた S_i の値を求めることができれば十分。


上記の課題を全体で O(DN) で解決することを考える。

 i を固定して考える。

このとき、 S_i は「元々のグラフのうち、辺の両端の持つ値がともに i で割り切れる辺だけを結合して得られる森において、同一の木に所属する相違なる 2 点間のパス上に存在する点の個数の総和」

と言い換えることができる。

いま、木のうち頂点が 1 つだけから構成されるものは答えに影響を与えないため無視することにすれば、 E_i を「元々のグラフのうち辺の両端の持つ値がともに iで割り切れる辺の個数」としてABC220F - Distance Sums 2の応用で O(E_i) で求めることができる。

また、 E_i の合計は「各辺の両端の持つ値の最大公約数の約数の個数の総和」として抑えることができるため、 前計算によりそれぞれの i に対して両端が共に i で割り切れる辺の情報を配列に保管しておくことで、全体を通し更新は O(DN) で行える。


これらを整理し、適切な実装を行うと、時間計算量・空間計算量ともに O(DN) で解くことができる。

[C++による実装例]

#pragma GCC target("avx2")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")


#include<bits/stdc++.h>

#define ll long long
#define endl '\n'


/*
    math/DS Template
*/
template <class T = int>T gcd(T a, T b){return (b == 0) ? a : gcd(b, a % b);}
template <class T = int>T extgcd(T a,T b,T &x,T &y){T g = a;x = 1;y = 0;if (b != 0) {g = extgcd(b, a % b, y, x), y -= (a / b) * x;}return g;}
template<class T = int> T invMod(T a,T m){T x,y;if (extgcd(a, m, x, y) == 1) {return (x + m) % m;}else{return -1;}}


const int maxn = 1e5 + 10;
constexpr ll mod = 998244353;
using Pii  = std::pair<int,int>;
using namespace std;


//global variables
vector<vector<int>> divisors(maxn), graph(maxn);
vector<ll> subsize(maxn,1), dist(maxn), total_weights(maxn);
vector<bool> visited(maxn, false);
vector<vector<Pii>> connected_edges(maxn);


void divisor_init(){
    // divisor setup
    for (int d = 1; d < maxn; d++){
        for (int j = d; j < maxn; j += d){
            divisors[j].emplace_back(d);
        }
    }
    return;
}


ll distance_sums(const vector<vector<int>> & g, const vector<int> & possible_roots){
    //ref: https://atcoder.jp/contests/abc220/editorial/2693
    ll total_distance = 0;
    //first dfs
    for (auto & root : possible_roots){
        if (visited[root]){continue;}
        function<void(int, int, int)> dfs = [&](int p, int u, int d){
            visited[u] = true;
            dist[root] += (d == 0 ? d : d + 1);
            for (auto v : g[u]){
                if (v == p){continue;}
                dfs(u, v, d + 1);
                subsize[u] += subsize[v];
            }
        }; 
        dfs(-1, root, 0);
    }
    //second dfs
    for (auto & root : possible_roots){
        if (!visited[root]){continue;}
        int subsizetree_size = subsize[root];
        function<void(int, int)> dfs1 = [&](int p, int u){
            total_distance += dist[u];
            for (auto v : g[u]){
                if (v == p){continue;}
             dist[v] = dist[u] - 2 * subsize[v] + subsizetree_size;
                dfs1(u, v);
            }
            //reset all trace to reuse global variables
            subsize[u] = 1; dist[u] = 0;
            visited[u] = false;
        };
        dfs1(-1, root);
    }
    // counted twice
    return (total_distance / 2) % mod;
}


void slv(){
    int N; cin >> N;
    vector<int> A(N);
    for (int i = 0; i < N; i++){cin >> A[i];}
    for (int i = 0; i < N - 1; i++){
        int u,v;
        cin >> u >> v;
        u--; v--;
        int gcdval = gcd(A[u], A[v]);
        for (auto d : divisors[gcdval]){
            connected_edges[d].emplace_back(make_pair(u, v));
        }
    }

    ll answer = 0;

    for (int i = maxn - 1; i > 0; i--){
        ll total_weight = 0;
        vector<int> roots_candidates;
        for (auto & [u, v] : connected_edges[i]){
            graph[u].emplace_back(v);
            graph[v].emplace_back(u);
            roots_candidates.emplace_back(u);
            roots_candidates.emplace_back(v);
        }
        //graph dynamically changes
        total_weight = distance_sums(graph, roots_candidates);
        //reset graph
        for (auto & root : roots_candidates){
            graph[root].clear();
        }
        // inclusion-exclusion principle
        for (int j = 2 * i; j < maxn; j += i){
            total_weight = (total_weight - total_weights[j]) % mod;
        }
        total_weights[i] = total_weight;
        answer += 1ll * i * total_weight; 
        answer %= mod;
    }

    if  (answer < 0){
        answer += mod;
    }

    cout << answer << endl;
    return;
}


signed main(){
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    divisor_init();
    int testcase = 1;
    while(testcase--){
        slv();
    }
    return 0;
}

感想

あまり書くことがなく、ビミョーな内容になってしまった...。

あとはてなブログmarkdown書きづらすぎる!