sotanishy's competitive programming library

sotanishy's code snippets for competitive programming

View the Project on GitHub sotanishy/cp-library-cpp

:heavy_check_mark: test/yosupo/rooted_tree_isomorphism_classification.hash.test.cpp

Depends on

Code

#define PROBLEM \
    "https://judge.yosupo.jp/problem/rooted_tree_isomorphism_classification"

#include <bits/stdc++.h>

#include "../../misc/compress.hpp"
#include "../../tree/tree_isomorphism.hpp"
using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    int N;
    cin >> N;
    vector<vector<int>> G(N);
    for (int i = 1; i < N; ++i) {
        int p;
        cin >> p;
        G[p].push_back(i);
    }
    TreeHasher hash;
    auto val = hash.hash_subtrees(G, 0);
    Compress<long long> comp(val);
    auto val2 = comp.compress(val);
    cout << comp.size() << endl;
    for (int i = 0; i < N; ++i) cout << val2[i] << (i < N - 1 ? " " : "\n");
}
#line 1 "test/yosupo/rooted_tree_isomorphism_classification.hash.test.cpp"
#define PROBLEM \
    "https://judge.yosupo.jp/problem/rooted_tree_isomorphism_classification"

#include <bits/stdc++.h>

#line 4 "misc/compress.hpp"

/**
 * @brief Coordinate Compression
 */
template <typename T>
class Compress {
   public:
    Compress() = default;
    explicit Compress(const std::vector<T>& vs) : xs(vs) {
        std::ranges::sort(xs);
        xs.erase(std::ranges::unique(xs).begin(), xs.end());
    }

    int compress(const T& x) const {
        return std::ranges::lower_bound(xs, x) - xs.begin();
    }

    std::vector<int> compress(const std::vector<T>& vs) const {
        std::vector<int> res(vs.size());
        std::ranges::transform(vs, res.begin(),
                               [&](const T& x) { return compress(x); });
        return res;
    }

    T decompress(int i) const { return xs[i]; }

    int size() const { return xs.size(); }

   private:
    std::vector<T> xs;
};
#line 7 "tree/tree_isomorphism.hpp"

#line 5 "tree/tree_diameter.hpp"

std::pair<int, std::vector<int>> tree_diameter(
    const std::vector<std::vector<int>>& G) {
    std::vector<int> to(G.size());

    auto dfs = [&](const auto& dfs, int v, int p) -> std::pair<int, int> {
        std::pair<int, int> ret(0, v);
        for (int c : G[v]) {
            if (c == p) continue;
            auto weight = dfs(dfs, c, v);
            ++weight.first;
            if (ret < weight) {
                ret = weight;
                to[v] = c;
            }
        }
        return ret;
    };

    auto p = dfs(dfs, 0, -1);
    auto q = dfs(dfs, p.second, -1);
    std::vector<int> path;
    int v = p.second;
    while (v != q.second) {
        path.push_back(v);
        v = to[v];
    }
    path.push_back(v);
    return {q.first, path};
}

template <typename T>
std::pair<T, std::vector<int>> tree_diameter(
    const std::vector<std::vector<std::pair<int, T>>>& G) {
    std::vector<int> to(G.size());

    auto dfs = [&](const auto& dfs, int v, int p) -> std::pair<T, int> {
        std::pair<T, int> ret(0, v);
        for (auto& [u, w] : G[v]) {
            if (u == p) continue;
            auto weight = dfs(dfs, u, v);
            weight.first += w;
            if (ret < weight) {
                ret = weight;
                to[v] = u;
            }
        }
        return ret;
    };

    auto p = dfs(dfs, 0, -1);
    auto q = dfs(dfs, p.second, -1);
    std::vector<int> path;
    int v = p.second;
    while (v != q.second) {
        path.push_back(v);
        v = to[v];
    }
    path.push_back(v);
    return {q.first, path};
}
#line 9 "tree/tree_isomorphism.hpp"

class TreeHasher {
   public:
    TreeHasher() : rng(rd()), rand(1, mod - 1) {}

    long long hash_all(const std::vector<std::vector<int>>& G, int root = -1) {
        long long res;
        if (root == -1) {
            auto [d, path] = tree_diameter(G);
            res = dfs_all(G, path[d / 2], -1).first;
            if (d % 2 == 1) {
                res = std::min(res, dfs_all(G, path[d / 2 + 1], -1).first);
            }
        } else {
            res = dfs_all(G, root, -1).first;
        }
        return res;
    }

    std::vector<long long> hash_subtrees(const std::vector<std::vector<int>>& G,
                                         int root) {
        std::vector<long long> hash(G.size());
        dfs_subtrees(G, hash, root, -1);
        return hash;
    }

   private:
    static constexpr long long mod = (1LL << 61) - 1;

    static inline long long add(long long a, long long b) {
        if ((a += b) >= mod) a -= mod;
        return a;
    }

    static inline long long mul(long long a, long long b) {
        __int128_t c = (__int128_t)a * b;
        return add(c >> 61, c & mod);
    }

    std::random_device rd;
    std::mt19937_64 rng;
    std::uniform_int_distribution<long long> rand;
    std::vector<long long> R;

    std::pair<long long, int> dfs_all(const std::vector<std::vector<int>>& G,
                                      int v, int p) {
        int maxd = 0;
        std::vector<long long> hash;
        for (int c : G[v]) {
            if (c != p) {
                auto [h, d] = dfs_all(G, c, v);
                maxd = std::max(maxd, d + 1);
                hash.push_back(h);
            }
        }
        if ((int)R.size() == maxd) {
            R.push_back(rand(rng));
        }
        long long res = 1;
        for (auto h : hash) {
            res = mul(res, add(R[maxd], h));
        }
        return {res, maxd};
    }

    int dfs_subtrees(const std::vector<std::vector<int>>& G,
                     std::vector<long long>& hash, int v, int p) {
        int maxd = 0;
        for (int c : G[v]) {
            if (c != p) {
                maxd = std::max(maxd, dfs_subtrees(G, hash, c, v) + 1);
            }
        }
        if ((int)R.size() == maxd) {
            R.push_back(rand(rng));
        }
        long long res = 1;
        for (int c : G[v]) {
            if (c != p) {
                res = mul(res, add(R[maxd], hash[c]));
            }
        }
        hash[v] = res;
        return maxd;
    }
};

class TreeEncoder {
   public:
    TreeEncoder() { mp[{}] = 0; }

    std::vector<int> encode(const std::vector<std::vector<int>>& G, int root) {
        std::vector<int> val(G.size());
        dfs(G, val, root, -1);
        return val;
    }

   private:
    std::map<std::vector<int>, int> mp;
    std::vector<long long> R;

    void dfs(const std::vector<std::vector<int>>& G, std::vector<int>& val,
             int v, int p) {
        std::vector<int> ch;
        for (int c : G[v]) {
            if (c != p) {
                dfs(G, val, c, v);
                ch.push_back(val[c]);
            }
        }
        std::ranges::sort(ch);
        if (!mp.contains(ch)) {
            mp[ch] = mp.size();
        }
        val[v] = mp[ch];
    }
};
#line 8 "test/yosupo/rooted_tree_isomorphism_classification.hash.test.cpp"
using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    int N;
    cin >> N;
    vector<vector<int>> G(N);
    for (int i = 1; i < N; ++i) {
        int p;
        cin >> p;
        G[p].push_back(i);
    }
    TreeHasher hash;
    auto val = hash.hash_subtrees(G, 0);
    Compress<long long> comp(val);
    auto val2 = comp.compress(val);
    cout << comp.size() << endl;
    for (int i = 0; i < N; ++i) cout << val2[i] << (i < N - 1 ? " " : "\n");
}
Back to top page