题解:P13875 [蓝桥杯 2024 省研究生组] 植物生命力

· · 题解

题目传送门

形式化题意

给定有根树 T,根为 s。设 S_i 为以 i 为根的子树集合,求

\sum_{i=1}^n\sum_{j\in S_i\setminus \{i\}}(![a_j<a_i])\times (![a_j\mid a_i])

题解

如果我们能把每个 i 对应的桶集合 B_i=\{a_j|j\in S_i\setminus\{i\}\} 快速的求解出来,那么我们就可以将每次判定转化为如下问题:求 j\in S_i\setminus \{i\} 内所有 a_j<a_ij 的数量减去所有 a_j<a_ia_j\mid a_i 的数量。前者可以使用数据结构维护 B_i 进行 O(\log n) 的单次求解前缀和,后者可以 O(\sqrt n) 枚举约数并查询是否存在于 B_i 进行求解。也就是,在不考虑快速求解 B_i 的情况下,我们已经可以做到 O(n\sqrt n+n\log n) 的时间复杂度。

那么如何求解 B_i?可以考虑合并儿子的 B_i,即 B_i=\cup_{j\in son_i}B_j。可以运用线段树合并或者 dsu on tree 进行维护,枚举约数的时间复杂度不受影响,可以做到 O(n\sqrt n+n\log n) 或者 O(n\sqrt n + n\log^2 n)

当然你也可以运用调和级数预处理约数,此时再合并 B_i 的时候套用线段树合并即可做到 O(n\log n)。当然由于枚举约数 + dsu on tree 也能通过并且码量较小,所以我选择实现这个。

实现(枚举约数 + dsu on tree)

#include <bits/stdc++.h>
using namespace std;
#define il inline
#define gc getchar
#define pb push_back
typedef long long ll;
const int N = 1e5 + 10;
int a[N], siz[N], dfn[N], rdfn[N], dp[N], son[N], b[N], c[N], dfscnt, n;
vector<int> g[N];
il void upd(int pos, int x)
{
    if(!pos) return;
    for(;pos <= n;pos += (pos & -pos)) c[pos] += x;
}
il int ask(int pos)
{
    int ans = 0;
    for(;pos > 0;pos -= (pos & -pos)) ans += c[pos];
    return ans;
}
il void dfs(int u, int fa)
{
    son[u] = -1;
    siz[u] = 1;
    dfn[u] = dp[u] = ++dfscnt;
    rdfn[dfscnt] = u;
    for(auto v : g[u])
    {
        if(v == fa) continue;
        dfs(v, u);
        siz[u] += siz[v];
        if(son[u] == -1 || siz[v] > siz[son[u]]) son[u] = v;
        dp[u] = max(dp[u], dp[v]);
    }
}
ll ans;
il void dfs2(int u, int fa, bool tag)
{
    for(auto v : g[u])
        if(v != fa && v != son[u]) dfs2(v, u, 1);
    if(son[u] != -1) dfs2(son[u], u, 0);
    for(auto v : g[u])
    {
        if(v == fa || v == son[u]) continue;
        for(int j = dfn[v];j <= dp[v];++j) b[a[rdfn[j]]]++, upd(a[rdfn[j]], 1);
    }
    ll cnt = 0;
    for(int i = 1;i <= a[u] / i;++i)
    {
        if(a[u] % i == 0)
        {
            if(i < a[u]) cnt += b[i];
            if(a[u] / i != i && a[u] / i < a[u]) cnt += b[a[u] / i];
        }
    }
    ans += ask(a[u] - 1) - cnt;
    upd(a[u], 1);
    b[a[u]]++;
    if(tag)
    {
        for(int i = dfn[u];i <= dp[u];++i) b[a[rdfn[i]]]--, upd(a[rdfn[i]], -1);
    }
}
il int read()
{
    char ch = gc();
    while(ch < 48 || ch > 57) ch = gc();
    int x = 0;
    while(48 <= ch && ch <= 57)
    {
        x = (x << 1) + (x << 3) + ch - 48;
        ch = gc();
    }
    return x;
}
int main()
{
    n = read();
    int s = read();
    for(int i = 1;i <= n;++i) a[i] = read();
    for(int i = 1;i < n;++i)
    {
        int u = read(), v = read();
        g[u].pb(v), g[v].pb(u);
    }
    dfs(s, 0);
    dfs2(s, 0, 0);
    cout << ans;
    return 0;
}