SPOJ COT Count on a tree(主席树+LCA)解题报告

2017-09-24 526 次浏览 解题报告

解题思路

主席树的区间第 K 大的树上版本。

从根结点到每个结点都建立一个版本的线段树即可。查询时 a 到 b 路径时需要按照 $root[a] + root[b] - root[lca(a,b)] - root[fa[lca(a,b)]]$ 计算。

参考代码

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <queue>
using namespace std;

const int MAXN = 100001;
const int DEG = 20;

struct edge {
    int to, next;
} e[MAXN<<1];

struct node {
    int sum, l, r;
} hjt[MAXN*40];

int a[MAXN], sorted[MAXN], num;
bool is_son[MAXN];    // 是否为子结点,唯一一个不为 true 的即是树根
int fa[MAXN][DEG];    // fa[i][j]表示结点i的第2^j个祖先
int deg[MAXN];    // 深度数组
int root[MAXN], cnt;
int head[MAXN], tot;

void AddEdge(int u, int v) {
    e[tot].to = v;
    e[tot].next = head[u];
    head[u] = tot++;
}

int GetIdx(int v) {
    return lower_bound(sorted+1, sorted+1+num, v) - sorted;
}

void Init() {
    tot = 0;
    memset(head, -1, sizeof(head));
    cnt = 0;
    root[0] = 0;
}

inline int CreateNode(int sum, int l, int r) {
    int idx = ++cnt;
    hjt[idx].sum = sum;
    hjt[idx].l = l;
    hjt[idx].r = r;

    return idx;
}

// 新建一棵线段树,只沿更新路径新建出较上个版本有修改的结点
// 调用参数
// root: 插入后新生成的线段树的根结点会赋值到 root 中存储
// pre_rt: 上一棵线段树的根
// pos: 本次要插入的数在线段树中的位置
// l, r: 递归参数。默认填写 1, num
void Insert(int &root, int pre_rt, int pos, int l, int r) {
    root = CreateNode(hjt[pre_rt].sum+1, hjt[pre_rt].l, hjt[pre_rt].r);
    if(l == r) return;
    int m = (l+r) >> 1;
    if(pos <= m)
        Insert(hjt[root].l, hjt[pre_rt].l, pos, l, m);
    else Insert(hjt[root].r, hjt[pre_rt].r, pos, m+1, r);
}

void BFS(int r) {
    queue<int> q;
    deg[r] = 0;
    fa[r][0] = r;
    q.push(r);
    Insert(root[r], root[0], GetIdx(a[r]), 1, num);
    while(!q.empty()) {
        int f = q.front();
        q.pop();
        for(int i=1; i<DEG; ++i) {
            fa[f][i] = fa[fa[f][i-1]][i-1];
        }
        for(int i=head[f]; ~i; i=e[i].next) {
            int v = e[i].to;
            if(v == fa[f][0]) continue;
            deg[v] = deg[f] + 1;
            fa[v][0] = f;
            q.push(v);
            Insert(root[v], root[f], GetIdx(a[v]), 1, num);
        }
    }
}

// 倍增 LCA
int LCA(int u, int v) {
    if(deg[u] > deg[v]) swap(u,v);
    int hu = deg[u], hv = deg[v];
    int tu = u, tv = v;
    for(int det=hv-hu, i=0; det ; det>>=1, ++i) {
        if(det&1) tv = fa[tv][i];
    }
    if(tu == tv) return tu;
    for(int i=DEG-1; i>=0; --i) {
        if(fa[tu][i] == fa[tv][i]) continue;
        tu = fa[tu][i];
        tv = fa[tv][i];
    }

    return fa[tu][0];
}

// 本函数适用于查询树上路径 (u, v) 中的第 k 小。
// 调用参数
// tu, tv, ta, tfa: u 的根,v 的根,LCA(u, v) 的根,LCA(u, v) 的父结点的根
// k: 要查询区间第几小
// l, r: 递归参数。默认填写 1, num
int Query(int tu, int tv, int ta, int tfa, int k, int l, int r) {
    if(l == r) return l;
    int m = (l+r) >> 1;
    int sum = hjt[hjt[tu].l].sum + hjt[hjt[tv].l].sum - hjt[hjt[ta].l].sum - hjt[hjt[tfa].l].sum;
    if(k <= sum)
        return Query(hjt[tu].l, hjt[tv].l, hjt[ta].l, hjt[tfa].l, k, l, m);
    else return Query(hjt[tu].r, hjt[tv].r, hjt[ta].r, hjt[tfa].r, k-sum, m+1, r);
}

int main(int argc, char const *argv[]) {
    int n, m, u, v, k;
    while(~ scanf("%d %d", &n, &m)) {
        Init();
        for(int i=1; i<=n; ++i) {
            scanf("%d", &a[i]);
            sorted[i] = a[i];
        }
        sort(sorted+1, sorted+1+n);
        num = unique(sorted+1, sorted+1+n) - (sorted+1);
        for(int i=0; i<n-1; ++i) {
            scanf("%d %d", &u, &v);
            AddEdge(u, v);
            AddEdge(v, u);
            is_son[v] = true;
        }
        // 查找树的根结点
        int r;
        for(int i=1; i<=n; ++i) {
            if(!is_son[i]) {
                r = i;
                break;
            }
        }
        BFS(r);
        fa[r][0] = 0;    // 令根结点的父结点为 0,防止 Query 时重复计算导致错误
        while(m--) {
            scanf("%d %d %d", &u, &v, &k);
            printf("%d\n", sorted[Query(root[u], root[v], root[LCA(u, v)], root[fa[LCA(u, v)][0]], k, 1, num)]);
        }
    }
    
    return 0;
}

bLue 创作,采用 知识共享署名 3.0,可自由转载、引用,但需署名作者且注明文章出处。
本文地址:https://dreamer.blue/blog/post/2017/09/24/spoj-cot.dream

还不快抢沙发

添加新评论