ZOJ red black tree

Posted 2021-01-07 aragaki

#include <cstdio>
#include <cstring>
#include <cmath>
#include <queue>
#include <map>
#include <vector>
#include <iostream>
#include <algorithm>
using namespace std;
typedef long long ll;
const int maxn = 100050;
const ll INF = (1LL << 62) - 1;
const double eps = 1e-8;

int t, n, m, q, k, x, no, a, b, pre[maxn][35];
int head[maxn], dep[maxn], nred[maxn], p[maxn];
ll dis[maxn], w, ans, res;
bool vis[maxn], flag[maxn];
struct node
{
        int to, nxt;
        ll w;
} e[maxn << 1];

void add(int a, int b, ll w)
{
        e[no].to = b, e[no].nxt = head[a], e[no].w = w;
        head[a] = no++;
        e[no].to = a, e[no].nxt = head[b], e[no].w = w;
        head[b] = no++;
}

bool cmp(int a, int b)
{
        return dis[a] - dis[nred[a]] > dis[b] - dis[nred[b]];
}

void init_lca()
{
        for (int j = 1; (1 << j) <= n; j++)
        {
                for (int i = 1; i <= n; i++)
                {
                        pre[i][j] = -1;
                }
        }
        for (int j = 1; (1 << j) <= n; j++)
        {
                for (int i = 1; i <= n; i++)
                {
                        if (pre[i][j - 1] != -1)
                        {
                                pre[i][j] = pre[pre[i][j - 1]][j - 1];
                        }
                }
        }
}

int lca(int x, int y)
{
        if (dep[x] < dep[y])
        {
                swap(x, y);
        }
        int mlg = 0;
        while ((1 << mlg) <= dep[x])
        {
                mlg++;
        }
        mlg--;
        for (int i = mlg; i >= 0; i--)
        {
                if (dep[x] - (1 << i) >= dep[y])
                {
                        x = pre[x][i];
                }
        }
        if (x == y)
        {
                return x;
        }
        for (int i = mlg; i >= 0; i--)
        {
                if (pre[x][i] != -1 && pre[x][i] != pre[y][i])
                {
                        x = pre[x][i], y = pre[y][i];
                }
        }
        return pre[x][0];
}

void init()
{
        no = dis[1] = 0;
        memset(vis, 0, sizeof(vis));
        memset(head, -1, sizeof(head));
        memset(flag, 0, sizeof(flag));
}

void dfs(int u, int red, int fa, int num)
{
        flag[u] = 1;
        nred[u] = red;
        dep[u] = num;
        pre[u][0] = fa;
        for (int i = head[u]; i != -1; i = e[i].nxt)
        {
                int v = e[i].to;
                if (flag[v])
                {
                        continue;
                }
                dis[v] = dis[u] + e[i].w;
                if (vis[v])
                {
                        dfs(v, v, u, num + 1);
                }
                else
                {
                        dfs(v, red, u, num + 1);
                }
        }
}

int main()
{
        scanf("%d", &t);
        while (t--)
        {
                scanf("%d%d%d", &n, &m, &q);
                init();
                for (int i = 1; i <= m; i++)
                {
                        scanf("%d", &x);
                        vis[x] = 1;
                }
                for (int i = 1; i < n; i++)
                {
                        scanf("%d%d%lld", &a, &b, &w);
                        add(a, b, w);
                }
                dfs(1, 1, 0, 1);
                init_lca();
                while (q--)
                {
                        scanf("%d", &k);
                        for (int i = 0; i < k; i++)
                        {
                                scanf("%d", &p[i]);
                        }
                        sort(p, p + k, cmp);
                        int now = p[0], i = 0;
                        ans = dis[p[0]] - dis[nred[p[0]]];
                        while (i < k)
                        {
                                int Lca = lca(now, p[i]);
                                if (nred[Lca] != nred[p[0]])
                                {
                                        break;
                                }
                                while (i < k && Lca == lca(now, p[i]) && nred[now] == nred[p[i]])
                                {
                                        i++;
                                }
                                res = dis[p[0]] - dis[Lca]; //初始最大值点到LCA的值
                                if (i < k)
                                {
                                        res =  max(res, dis[p[i]] - dis[nred[p[i]]]); //LCA子树之外的最大值
                                }
                                if (res < ans)
                                {
                                        ans = res;
                                }
                                else
                                {
                                        break;
                                }
                                now = Lca;
                        }
                        printf("%lld
", ans);
                }
        }
        return 0;
}

View Code