test/aoj0439.test.cpp

• View this file on GitHub
• Last update: 2021-06-21 15:24:20+09:00
• Problem: http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=0439

Depends on

• datastructure/sparsetable.cpp
• tree/auxtree.cpp

Code

#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=0439"
#include <iostream>
#include <algorithm>
#include <map>
#include <set>
#include <queue>
#include <stack>
#include <numeric>
#include <bitset>
#include <cmath>

static const int MOD = 1000000007;
using ll = long long;
using uint = unsigned;
using ull = unsigned long long;
using namespace std;

template<class T> constexpr T INF = ::numeric_limits<T>::max()/32*15+208;

#include "../tree/auxtree.cpp"

template <typename T>
using GPQ = priority_queue<T, vector<T>, greater<T>>;
int main() {
    int n;
    cin >> n;
    vector<vector<int>> color(n);
    AuxTree G(n);
    for (int i = 0; i < n; ++i) {
        int x;
        scanf("%d", &x);
        color[x-1].emplace_back(i);
    }
    for (int i = 0; i < n-1; ++i) {
        int u, v;
        scanf("%d %d", &u, &v);
        u--; v--;
        G.add_edge(u, v);
    }
    G.buildLCA();
    vector<int> ans(n, INF<int>);
    vector<int> cmp(n), dist(n);
    for (auto &&v : color) {
        if(v.empty()) continue;
        int k = v.size();
        G.make(v);
        GPQ<pair<int, int>> Q;
        for (int i = 0; i < v.size(); ++i) {
            if(i < k) dist[v[i]] = 0, cmp[v[i]] = v[i], Q.emplace(0, v[i]);
            else dist[v[i]] = INF<int>;
        }
        while(!Q.empty()){
            auto [d, i] = Q.top(); Q.pop();
            if(d != dist[i]) continue;
            for (auto &&j : G.out[i]) {
                if(dist[j] > dist[i] + G.distance(i, j)){
                    dist[j] = dist[i] + G.distance(i, j);
                    cmp[j] = cmp[i];
                    Q.emplace(dist[j], j);
                }
            }
        }
        for (auto &&i : v) {
            for (auto &&j : G.out[i]) {
                int l = cmp[i], r = cmp[j];
                if(l != r) ans[l] = min(ans[l], G.distance(l, r));
            }
        }
        G.clear(v);
    }
    for (auto &&i : ans) {
        printf("%d\n", i);
    }
    return 0;
}

#line 1 "test/aoj0439.test.cpp"
#define PROBLEM "http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=0439"
#include <iostream>
#include <algorithm>
#include <map>
#include <set>
#include <queue>
#include <stack>
#include <numeric>
#include <bitset>
#include <cmath>

static const int MOD = 1000000007;
using ll = long long;
using uint = unsigned;
using ull = unsigned long long;
using namespace std;

template<class T> constexpr T INF = ::numeric_limits<T>::max()/32*15+208;

#line 1 "datastructure/sparsetable.cpp"
template <class F>
struct SparseTable {
    using T = typename F::T;
    vector<vector<T>> table;
    vector<int> u;
    SparseTable() = default;
    explicit SparseTable(const vector<T> &v){ build(v); }
 
    void build(const vector<T> &v){
        int n = v.size(), m = 1;
        while((1<<m) <= n) m++;
        table.assign(m, vector<T>(n));
        u.assign(n+1, 0);
        for (int i = 2; i <= n; ++i) {
            u[i] = u[i>>1] + 1;
        }
        for (int i = 0; i < n; ++i) {
            table[0][i] = v[i];
        }
        for (int i = 1; i < m; ++i) {
            int x = (1<<(i-1));
            for (int j = 0; j < n; ++j) {
                table[i][j] = F::f(table[i-1][j], table[i-1][min(j+x, n-1)]);
            }
        }
    }
 
    T query(int a, int b){
        int l = b-a;
        return F::f(table[u[l]][a], table[u[l]][b-(1<<u[l])]);
    }
};
#line 2 "tree/auxtree.cpp"

struct F {
    using T = pair<int, int>;
    static T f(T a, T b) { return min(a, b); }
    static T e() { return T{INF<int>, -1}; }
};

class AuxTree {
    SparseTable<F> table;
    void dfs_euler(int v, int p, int d, int &k, int &l){
        id[v] = k;
        vs[k] = v;
        depth[k++] = d;
        dep[v] = d;
        fi[v] = l++;
        for (auto &&u : G[v]) {
            if(u != p){
                dfs_euler(u, v, d+1, k, l);
                vs[k] = v;
                depth[k++] = d;
            }
        }
    }
public:
    int n;
    vector<vector<int>> G, out;
    vector<int> vs, depth, dep, id, fi;
    explicit AuxTree(int n) : n(n), G(n), out(n), vs(2*n-1), depth(2*n-1), dep(n), id(n), fi(n), table() {};
    void add_edge(int a, int b){
        G[a].emplace_back(b);
        G[b].emplace_back(a);
    }

    void eulertour(int root) {
        int k = 0, l = 0;
        dfs_euler(root, -1, 0, k, l);
    }

    void buildLCA(){
        eulertour(0);
        vector<pair<int, int>> v(2*n-1);
        for (int i = 0; i < 2*n-1; ++i) {
            v[i] = make_pair(depth[i], vs[i]);
        }
        table.build(v);
    }

    void make(vector<int> &v){
        sort(v.begin(),v.end(), [&](int a, int b){ return fi[a] < fi[b]; });
        v.erase(unique(v.begin(), v.end()), v.end());
        int k = v.size();
        stack<int> s;
        s.emplace(v.front());
        for (int i = 0; i+1 < k; ++i) {
            int w = LCA(v[i], v[i+1]);
            if(w != v[i]){
                int u = s.top(); s.pop();
                while(!s.empty() && dep[w] < dep[s.top()]){
                    out[s.top()].emplace_back(u);
                    out[u].emplace_back(s.top());
                    u = s.top(); s.pop();
                }
                if(s.empty() || s.top() != w){
                    s.emplace(w);
                    v.emplace_back(w);
                }
                out[w].emplace_back(u);
                out[u].emplace_back(w);
            }
            s.emplace(v[i+1]);
        }
        while(s.size() > 1){
            int u = s.top(); s.pop();
            out[s.top()].emplace_back(u);
            out[u].emplace_back(s.top());
        }
    }

    void clear(vector<int> &v){
        for (auto &&i : v) {
            out[i].clear();
            out[i].shrink_to_fit();
        }
    }

    int LCA(int u, int v){
        if(id[u] > id[v]) swap(u, v);
        return table.query(id[u], id[v]+1).second;
    }

    int distance(int u, int v){
        return dep[u]+dep[v]-2*dep[LCA(u, v)];
    }
};
#line 21 "test/aoj0439.test.cpp"

template <typename T>
using GPQ = priority_queue<T, vector<T>, greater<T>>;
int main() {
    int n;
    cin >> n;
    vector<vector<int>> color(n);
    AuxTree G(n);
    for (int i = 0; i < n; ++i) {
        int x;
        scanf("%d", &x);
        color[x-1].emplace_back(i);
    }
    for (int i = 0; i < n-1; ++i) {
        int u, v;
        scanf("%d %d", &u, &v);
        u--; v--;
        G.add_edge(u, v);
    }
    G.buildLCA();
    vector<int> ans(n, INF<int>);
    vector<int> cmp(n), dist(n);
    for (auto &&v : color) {
        if(v.empty()) continue;
        int k = v.size();
        G.make(v);
        GPQ<pair<int, int>> Q;
        for (int i = 0; i < v.size(); ++i) {
            if(i < k) dist[v[i]] = 0, cmp[v[i]] = v[i], Q.emplace(0, v[i]);
            else dist[v[i]] = INF<int>;
        }
        while(!Q.empty()){
            auto [d, i] = Q.top(); Q.pop();
            if(d != dist[i]) continue;
            for (auto &&j : G.out[i]) {
                if(dist[j] > dist[i] + G.distance(i, j)){
                    dist[j] = dist[i] + G.distance(i, j);
                    cmp[j] = cmp[i];
                    Q.emplace(dist[j], j);
                }
            }
        }
        for (auto &&i : v) {
            for (auto &&j : G.out[i]) {
                int l = cmp[i], r = cmp[j];
                if(l != r) ans[l] = min(ans[l], G.distance(l, r));
            }
        }
        G.clear(v);
    }
    for (auto &&i : ans) {
        printf("%d\n", i);
    }
    return 0;
}

Back to top page