CodeChef July 18 SUBWAY

【模板】树的点分治.cpp

#include <bits/stdc++.h>

using namespace std;
#define pb push_back
#define eb emplace_back
#define all(x) x.begin(), x.end()
#define debug(x) cerr << #x <<": " << (x) << endl
#define DEBUG printf("Passing [%s] in LINE %d\n",__FUNCTION__,__LINE__)
#ifdef LOCAL
#define see(x) cout  << #x << ": " << (x) << endl
#endif
#ifndef LOCAL
#define see(x)
#endif


#define rep(n) for(int _ = 0; _ != (n); ++_)
//#define rep(i, a, b) for(int i = (a); i <= (b); ++i)
#define Rng(i, n) for(int i = 0; i != (n); ++i)
#define rng(i, a, b) for(int i = (a); i < (b); ++i)
#define RNG(i, a) for(auto &i: (a))
#define dwn(i, r, l) for(int i = (r); i>=(l); i--)

namespace std {
    template<class T>
    T begin(std::pair<T, T> p)
    {
        return p.first;
    }
    template<class T>
    T end(std::pair<T, T> p)
    {
        return p.second;
    }
}


#if __cplusplus < 201402L
template<class Iterator>
std::reverse_iterator<Iterator> make_reverse_iterator(Iterator it)
{
    return std::reverse_iterator<Iterator>(it);
}
#endif

template<class Range>
std::pair<std::reverse_iterator<decltype(begin(std::declval<Range>()))>, std::reverse_iterator<decltype(begin(std::declval<Range>()))>> make_reverse_range(Range &&r)
{
    return std::make_pair(make_reverse_iterator(::begin(r)), make_reverse_iterator(::end(r)));
}

#define RRNG(x, cont) for (auto &x: make_reverse_range(cont))



template<class T> int sign(const T &a) { return a == 0 ? 0 : a > 0 ? 1 : -1; }
template<class T> inline T min(T a, T b, T c){return min(min(a, b), c);}
template<class T> inline T max(T a, T b, T c){return max(max(a, b), c);}
template<class T> void Min(T &a, const T &b){ a = min(a, b); }
template<class T> void Max(T &a, const T &b){ a = max(a, b); }

template<typename T> void println(const T &t) { cout << t << '\n'; }
template<typename T, typename ...Args> void println(const T &t, const Args &...rest) { cout << t << ' '; println(rest...); }

template<typename T> void print(const T &t) { cout << t << ' '; }

template<typename T, typename ...Args> void print(const T &t, const Args &...rest) { cout << t; print(rest...); }

// this overload is chosen when there's only one argument
template<class T> void scan(T &t) { cin >> t; }
template<class T, class ...Args> void scan(T &a, Args &...rest) { cin >> a; scan(rest...); }

using ll = long long;
using ull = unsigned long long;
using vec = vector<ll>;
using mat = vector<vec>;
using pii = pair<int, int>;
using pdd = pair<double, double>;
using pip = pair<int, pii>;
using szt = size_t;
using vi = vector<int>;
using vl = vector<ll>;
using vb = vector<bool>;
using vpii = vector<pii>;
using vvi = vector<vi>;
using pli = pair<ll,int>;
using wg = vector<vpii>; //weighted graph

int cas;
const double pi = acos(-1);
ll mod = 1e9 + 7;

template<class T>
inline void add_mod(T &a, const T &b) {
    a += b;
    if (a >= mod) a -= mod;
}
template<class T>
void sub_mod(T &a, const T &b){
    a -= b;
    if (a < 0) a += mod;
}
auto bo=[](int x){
    bitset<5> a(x);
    cout << a << endl;
};



mat operator*(const mat &a, const mat &b) {
    mat c(a.size(), vec(b[0].size()));
    for (int i = 0; i < a.size(); i++) {
        for (int j = 0; j < a[0].size(); j++) {
            if (a[i][j]) { // optimization for sparse matrix
                for (int k = 0; k < b[0].size(); k++) {
                    add_mod(c[i][k], a[i][j] * b[j][k] % mod);
                }
            }
        }
    }
    return c;
}

vec operator*(const mat &a, const vec &b) {
    vec c(a.size());
    for (int i = 0; i < a.size(); i++) {
        for (int j = 0; j < a[0].size(); j++) {
            add_mod(c[i], a[i][j] * b[j] % mod);
        }
    }
    return c;
}

mat pow(mat a, ull n) {
    mat res(a.size(), vec(a[0].size()));
    for (int i = 0; i < a.size(); i++) {
        res[i][i] = 1;
    }
    while (n) {
        if (n & 1) {
            res = res * a;
        }
        a = a * a;
        n >>= 1;
    }
    return res;
}

std::ostream& operator<<(std::ostream& os, __int128 T) {
    if (T<0) os<<"-";
    if (T>=10 ) os<<T/10;
    if (T<=-10) os<<(-(T/10));
    return os<<( (int) (T%10) >0 ? (int) (T%10) : -(int) (T%10) ) ;
}

__int128 LPOW(__int128 x, ll n) {
    __int128 res = 1;
    for (; n; n /= 2, x *= x, x %= mod) {
        if (n & 1) {
            res *= x;
            res %= mod;
        }
    }
    return res;
}

ll POW(ll x, ll n){
    ll res = 1;
    for (; n; n /= 2, x *= x, x %= mod) {
        if (n & 1) {
            res *= x;
            res %= mod;
        }
    }
    return res;
}


ll INV(ll x) {
    return POW(x, mod - 2);
}

ll inv(ll x){
//    see(x);
    return x == 1? 1: (mod - mod/x * inv(mod%x) % mod);
}



// 2D rotation
void rotate(double &x, double &y, double theta) {
    double tx = cos(theta) * x - sin(theta) * y;
    double ty = sin(theta) * x + cos(theta) * y;
    x = tx, y = ty;
}
namespace bit {
    const int BIT_N = 1e5 + 5;

    int bit[BIT_N];

    int sum(int x) {
        int res = 0;
        while (x) {
            res += bit[x];
            x -= x & -x;
        }
        return res;
    }

    int sum(int l, int r) {
        if (l > r) return 0;
        return sum(r) - sum(l - 1);
    }

    void add(int x, int v, int n) {
        while (x <= n) {
            bit[x] += v;
            x += x & -x;
        }
    }
}

namespace util{
    int len(ll x){return snprintf(nullptr, 0, "%lld", x);}
    vi get_d(ll x){
        vi res;
        while(x) {
            res.pb(x%10);
            x /= 10;
        }
        reverse(all(res));
        return res;
    }
    template <class T> T parity(const T &a){
        return a & 1;
    }
    template <class T>
    void out (const vector<T> &a){
        std::copy(a.begin(), a.end(), std::ostream_iterator<T>(std::cout, ", "));
        cout << endl;
    };
}

using namespace util;

// #include <ext/pb_ds/priority_queue.hpp>

// typedef __gnu_pbds :: priority_queue<pip, less<pip>, __gnu_pbds::thin_heap_tag > Heap;

// Heap h;

// Heap::point_iterator pos[N][N];
const ll LINF = LLONG_MAX/10;
const int INF = INT_MAX/10;
const int M = 3000 + 5;


const int N = 5e5+5;

struct edge{
    int v, i, next;
} E[2*N];
int last[N];
int cnt;

void insert(int u, int v, int i){
    E[++cnt] = {v, i, last[u]}; last[u] = cnt;
    E[++cnt] = {u, i, last[v]}; last[v] = cnt;
}

bool used[N];
int size[N];
int ma[N];
int root, sum; // 当前子树中的节点数，求重心用

void centroid(int u, int f){
    size[u]=1, ma[u] = 0;
    for(int i = last[u]; i; i = E[i].next){
        int &v = E[i].v;
        if(v!=f && !used[v]){
            centroid(v, u);
            size[u]+=size[v];
            ma[u] = max(ma[u], size[v]);
        }
    }
    ma[u]=max(ma[u], sum-size[u]);
    if(ma[root] > ma[u]) root = u;
};

int last_edge2[N][2][2]; // 表示从u出发的最长路的末边颜色数是否唯一
inline int work(int u, int c) {
    if (c != last_edge2[u][0][0]){
        if (c != last_edge2[u][1][0]) return 0;
        else return last_edge2[u][0][1];
    } else return last_edge2[u][1][1];
};

// 重心分治T到死，启发式合并是正解？

int main() {
    // Single Cut of Failure taught me
    cout << std::fixed; //
    cout << setprecision(10);
//    ios::sync_with_stdio(false);
//    cin.tie(nullptr);
#ifdef LOCAL
    freopen("main.in", "r", stdin);
//    freopen("main.out", "w", stdout);
#endif
    int n, m;
    scanf("%d%d", &n, &m);
    wg color(n+1);
    vector<pip> col(n+1);
    int cnt = 0;
    rep(m){
        int u, v, w;
        scanf("%d%d%d", &u, &v, &w);
        if(u > v) swap(u, v);
        color[u].eb(v, w);
    }
    rng(i, 1, n+1){
        sort(all(color[i]));
        color[i].erase(unique(all(color[i])), color[i].end());
        int j = 0;
        while(j < color[i].size()){
            int v = color[i][j].first;
            ++cnt;
            insert(i, v, cnt);
            for(;j < color[i].size() && color[i][j].first == v;++j){
                if(col[cnt].first == 3) continue;
                ++col[cnt].first;
                if(col[cnt].first == 1){
                    col[cnt].second.first = color[i][j].second;
                }
                else if(col[cnt].first==2){
                    col[cnt].second.second = color[i][j].second;
                }
            }
        }
    }
    int q;
    scanf("%d", &q);
    wg Q(n+1);
    rng(i, 0, q){
        int u, v;
        scanf("%d%d", &u, &v);
        if(u != v) Q[u].eb(v, i);
    }


    vi last_edge(n+1);
    vi d(n+1);

    function<void(int,int,int,int)> dfs;
    vi first_edge(n+1);
    vpii label(n+1);
    dfs = [&](int u, int f, int top, int num) {
        label[u] = {root, num};
        last_edge[u] = top == -1 ? 0 : top;
        for(int i = last[u]; i; i = E[i].next) {
            int v = E[i].v;
            if(v == f || used[v]) continue;
            auto &s = col[E[i].i];
            if(u == root) {
                first_edge[v] = s.first > 1 ? 0 : s.second.first;
                d[v] = 0;
                num = v;
            }
            else {
                auto &t = first_edge[u];
                if(t == 0) {
                    first_edge[v] = s.first > 1 ? 0 : s.second.first;
                    d[v] = d[u] + 1;
                }
                else {
                    if(s.first > 2) {
                        d[v] = d[u] + 1;
                        first_edge[v] = 0;
                    }
                    else if(s.first == 1) {
                        d[v] = d[u] + (t != s.second.first);
                        first_edge[v] = s.second.first;
                    }
                    else{
                        int c1 = s.second.first, c2 = s.second.second;
                        d[v] = d[u] + 1;
                        if(c1 == t) first_edge[v] = c2;
                        else if(c2 == t) first_edge[v] = c1;
                        else first_edge[v] = 0;
                    }
                }
            }

            int ntop = top;
            if(top == -1) {
                if (s.first > 2) ntop = 0;
                else if (s.first == 2) {
                    int c1 = s.second.first, c2 = s.second.second;
                    if (u == root) {
                        last_edge2[v][0][0] = last_edge2[v][0][1] = c1;
                        last_edge2[v][1][0] = last_edge2[v][1][1] = c2;
                    } else {
                        int r1 = work(u, c1), r2 = work(u, c2);
                        if (!r1 && !r2) ntop = 0;
                        else {
                            last_edge2[v][0][0] = c1;
                            last_edge2[v][0][1] = r1;
                            last_edge2[v][1][0] = c2;
                            last_edge2[v][1][1] = r2;
                        }
                    }
                } else {
                    int c = s.second.first;
                    ntop = u == root ? c : work(u, c);
                }
            }
            dfs(v, u, ntop, num);
        }
    };

    vi ans(q);
    function<void(int,int)> calc;
    calc = [&](int u, int f){
        RNG(x, Q[u]){
            int v = x.first, i = x.second;
            if(label[v].first == root) {
                if (u == root) {
                    ans[i] = d[v];
                } else if (v == root) {
                    ans[i] = d[u];
                } else if (label[u].second != label[v].second) {
                    if (last_edge[v] == 0 || last_edge[u] == 0) {
                        ans[i] = d[u] + d[v] + 1;
                    } else {
                        ans[i] = d[u] + d[v] + (last_edge[u] != last_edge[v]);
                    }
                }
            }
        }
        for(int i = last[u]; i; i=E[i].next){
            int &v = E[i].v;
            if(!used[v] && v != f)
                calc(v, u);
        }
    };
    // 原来我来点分治都不会！

    function<void(int)> DC;
    DC = [&](int u){ // u是当前找到的重心
        used[u] = true;
        dfs(u, 0, -1, 0);
        calc(u, 0);
        for(int i=last[u]; i; i=E[i].next){
            int v = E[i].v;
            if(!used[v]){
                sum = size[v]; //
                root = 0;
                centroid(v, u);
                DC(root);
            }
        }
    };

    ma[0] = INF;
    sum = n; // 调用 centroid 之前必不可少的初始化！
    centroid(1, 1);
    DC(root);

    RNG(x, ans) println(x);
#ifdef LOCAL
    cerr << "Time elapsed: " << 1.0 * clock() / CLOCKS_PER_SEC << " s.\n";
#endif
    return 0;
}

关于点分治的写法，我参考了 hzwer

但是代码的第457行是错的。v 有可能是重心 u 的父亲，应当把当前考虑的子树的节点数作为函数 DC 的参数。

    function<void(int, int)> DC;
    DC = [&](int u, int tot){ // u是当前找到的重心，tot 是当前子树的大小
        used[u] = true;
        dfs(u, 0, -1, 0);
        calc(u, 0);
        for(int i=last[u]; i; i=E[i].next){
            int v = E[i].v;
            if(!used[v]){
                sum = size[v] > size[u] ? tot - size[u] : size[v];  //
                root = 0;
                centroid(v, u);
                DC(root);
            }
        }
    };

或者写成

    function<void(int)> DC;
    DC = [&](int u){ // u是当前找到的重心
        used[u] = true;
        dfs(u, 0, -1, 0);
        calc(u, 0);
        int tot = sum;
        for(int i=last[u]; i; i=E[i].next){
            int v = E[i].v;
            if(!used[v]){
                sum = size[v] > size[u] ? tot - size[v] : size[v]; //
                root = 0;
                centroid(v, u);
                DC(root);
            }
        }
    };

我查了一下网上关于树的点分治的题解，大部分人的代码在这个地方的处理都跟 457 行类似。这个地方是值得注意的。

关于如何判断询问中两个点的 LCA 是否是当前考虑的重心，我的判断方法是给当前子树中除了重心之外的每个点 u 打两个标记，一个标记是当前的重心，另一个标记 u 所在的子树的根是重心的哪个儿子。

还有一种办法是预先求出所有的重心，这些重心也构成一棵树，称作「点分树」，我认为称作「重心树」更贴切。
询问中的两个点的 LCA 是哪些重心？即两点在重心树中的 LCA。
借助重心树求 LCA 有两种比较高效的做法：

1. 把重心树上各点的深度按 DFS 序存到一个数组里，然后构建稀疏表求 RMQ
2. 在求重心的同时用 Tarjan 的离线 LCA 算法

关于代码如何写得高效的一点经验：
尽量减少函数的参数个数。