POJ 2114 Boatherds 树分治

题意：

给出一颗边带权的\(n\)个节点的树，问是否存在最短距离为\(k\)的点对。

分析：

最开始做的题是询问最短距离小于等于\(k\)的点对。

我第一反应是下面的方法一。

方法一：
先求出小于等于\(k\)的点对 和 小于\(k\)(也就是小于等于\(k-1\))的点对，然后相减得到等于\(k\)的点对的个数。

方法二：
直接修改之前的统计点对的函数。

容易知道，第一种方法每组查询跑了两遍，所以运行时间是方法二的两倍。

方法一的代码：

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
#include <map>
#define MP make_pair
using namespace std;

void read(int& x) {
    x = 0;
    char c = ' ';
    while(c < '0' || c > '9') c = getchar();
    while('0' <= c && c <= '9') {
        x = x * 10 + c - '0';
        c = getchar();
    }
}

typedef pair<int, int> PII;
const int maxn = 10000 + 10;
const int INF = 0x3f3f3f3f;

struct Edge
{
    int v, w, nxt;
    Edge() {}
    Edge(int v, int w, int nxt): v(v), w(w), nxt(nxt) {}
};

int n, k, ans;
vector<int> d, d2;
bool del[maxn];

int ecnt, head[maxn];
Edge edges[maxn * 2];

void AddEdge(int u, int v, int w) {
    edges[ecnt] = Edge(v, w, head[u]);
    head[u] = ecnt++;
}

int fa[maxn], sz[maxn];

void dfs(int u) {
    sz[u] = 1;
    for(int i = head[u]; ~i; i = edges[i].nxt) {
        int v = edges[i].v;
        if(v == fa[u] || del[v]) continue;
        fa[v] = u;
        dfs(v);
        sz[u] += sz[v];
    }
}

PII findCenter(int u, int t) {
    PII ans(INF, -1);
    int m = 0;
    for(int i = head[u]; ~i; i = edges[i].nxt) {
        int v = edges[i].v;
        if(v == fa[u] || del[v]) continue;
        ans = min(ans, findCenter(v, t));
        m = max(m, sz[v]);
    }
    m = max(m, t - sz[u]);
    ans = min(ans, MP(m, u));
    return ans;
}

void getDist(int u, int p, int dist, vector<int>& d) {
    d.push_back(dist);
    for(int i = head[u]; ~i; i = edges[i].nxt) {
        int v = edges[i].v, w = edges[i].w;
        if(v == p || del[v]) continue;
        getDist(v, u, dist + w, d);
    }
}

int cntPiars(vector<int>& d) {
    int ans = 0;
    sort(d.begin(), d.end());
    int j = d.size();
    for(int i = 0; i < d.size(); i++) {
        while(j && d[i] + d[j-1] > k) j--;
        ans += j - (j > i ? 1 : 0);
    }
    return ans;
}

void solve(int u) {
    fa[u] = 0;
    dfs(u);
    int s = findCenter(u, sz[u]).second;
    del[s] = true;
    
    for(int i = head[u]; ~i; i = edges[i].nxt) {
        int v = edges[i].v;
        if(del[v]) continue;
        solve(v);
    }

    d.clear();
    d.push_back(0);
    for(int i = head[s]; ~i; i = edges[i].nxt) {
        int v = edges[i].v, w = edges[i].w;
        if(del[v]) continue;
        d2.clear();
        getDist(v, s, w, d2);
        ans -= cntPiars(d2);
        d.insert(d.begin(), d2.begin(), d2.end());
    }

    ans += cntPiars(d);
    del[s] = false;
}

int main()
{
    while(scanf("%d", &n) == 1 && n) {
        ecnt = 0;
        memset(head, -1, sizeof(head));
        for(int u = 1; u <= n; u++) {
            int v, w;
            for(read(v); v; read(v)) {
                read(w);
                AddEdge(u, v, w);
                AddEdge(v, u, w);
            }
        }

        for(read(k); k; read(k)) {
            ans = 0;
            solve(1);
            int pre = ans;
            ans = 0;
            k--;
            solve(1);
            printf("%s\n", pre - ans > 0 ? "AYE" : "NAY");
        }
        printf(".\n");
    }

    return 0;
}

方法二的代码：

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
#include <map>
#define MP make_pair
using namespace std;

void read(int& x) {
    x = 0;
    char c = ' ';
    while(c < '0' || c > '9') c = getchar();
    while('0' <= c && c <= '9') {
        x = x * 10 + c - '0';
        c = getchar();
    }
}

typedef pair<int, int> PII;
const int maxn = 10000 + 10;
const int INF = 0x3f3f3f3f;

struct Edge
{
    int v, w, nxt;
    Edge() {}
    Edge(int v, int w, int nxt): v(v), w(w), nxt(nxt) {}
};

int n, k, ans;
vector<int> d, d2;
bool del[maxn];

int ecnt, head[maxn];
Edge edges[maxn * 2];

void AddEdge(int u, int v, int w) {
    edges[ecnt] = Edge(v, w, head[u]);
    head[u] = ecnt++;
}

int fa[maxn], sz[maxn];

void dfs(int u) {
    sz[u] = 1;
    for(int i = head[u]; ~i; i = edges[i].nxt) {
        int v = edges[i].v;
        if(v == fa[u] || del[v]) continue;
        fa[v] = u;
        dfs(v);
        sz[u] += sz[v];
    }
}

PII findCenter(int u, int t) {
    PII ans(INF, -1);
    int m = 0;
    for(int i = head[u]; ~i; i = edges[i].nxt) {
        int v = edges[i].v;
        if(v == fa[u] || del[v]) continue;
        ans = min(ans, findCenter(v, t));
        m = max(m, sz[v]);
    }
    m = max(m, t - sz[u]);
    ans = min(ans, MP(m, u));
    return ans;
}

void getDist(int u, int p, int dist, vector<int>& d) {
    d.push_back(dist);
    for(int i = head[u]; ~i; i = edges[i].nxt) {
        int v = edges[i].v, w = edges[i].w;
        if(v == p || del[v]) continue;
        getDist(v, u, dist + w, d);
    }
}

int cntPiars(vector<int>& d) {
    int ans = 0;
    sort(d.begin(), d.end());
    int i = 0, j = d.size() - 1;
    while(i < j) {
        if(d[i] + d[j] < k) i++;
        else if(d[i] + d[j] > k) j--;
        else {
            if(d[i] == d[j]) {
                ans += (j - i + 1) * (j - i) / 2;
                break;
            } else {
                int l = i, r = j;
                while(d[i] == d[l]) i++;
                while(d[j] == d[r]) j--;
                ans += (i - l) * (r - j);
            }
        }
    }
    return ans;
}

void solve(int u) {
    fa[u] = 0;
    dfs(u);
    int s = findCenter(u, sz[u]).second;
    del[s] = true;
    
    for(int i = head[u]; ~i; i = edges[i].nxt) {
        int v = edges[i].v;
        if(del[v]) continue;
        solve(v);
    }

    d.clear();
    d.push_back(0);
    for(int i = head[s]; ~i; i = edges[i].nxt) {
        int v = edges[i].v, w = edges[i].w;
        if(del[v]) continue;
        d2.clear();
        getDist(v, s, w, d2);
        ans -= cntPiars(d2);
        d.insert(d.begin(), d2.begin(), d2.end());
    }

    ans += cntPiars(d);
    del[s] = false;
}

int main()
{
    while(scanf("%d", &n) == 1 && n) {
        ecnt = 0;
        memset(head, -1, sizeof(head));
        for(int u = 1; u <= n; u++) {
            int v, w;
            for(read(v); v; read(v)) {
                read(w);
                AddEdge(u, v, w);
                AddEdge(v, u, w);
            }
        }

        for(read(k); k; read(k)) {
            ans = 0;
            solve(1);
            printf("%s\n", ans ? "AYE" : "NAY");
        }
        printf(".\n");
    }

    return 0;
}