test/src/datastructure/union_find/union_find_component_sum/abc183_f.test.cpp

View this file on GitHub
Last update: 2026-06-19 20:35:33+09:00
Problem: https://atcoder.jp/contests/abc183/tasks/abc183_f

Depends on

Code

#define PROBLEM "https://atcoder.jp/contests/abc183/tasks/abc183_f"

#include <iostream>
#include <map>

#include "library/datastructure/union_find/union_find_component_sum.hpp"

using sum_type = std::map<int, int>;

void merge(sum_type &par_data, sum_type ch_data) {
    for (auto &[k, v] : ch_data) par_data[k] += v;
}

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int n, q;
    std::cin >> n >> q;

    std::vector<int> c(n);
    for (int i = 0; i < n; ++i) std::cin >> c[i];

    std::vector<sum_type> init_data(n);
    for (int i = 0; i < n; ++i) {
        ++init_data[i][c[i]];
    }
    suisen::UnionFindComponentSum<sum_type, merge> uf(init_data);

    while (q --> 0) {
        int query_type;
        std::cin >> query_type;
        if (query_type == 1) {
            int a, b;
            std::cin >> a >> b;
            --a, --b;
            uf.merge(a, b);
        } else {
            int x, y;
            std::cin >> x >> y;
            --x;
            const auto& sum = uf.sum(x);
            if (auto it = sum.find(y); it == sum.end()) {
                std::cout << 0 << '\n';
            } else {
                std::cout << it->second << '\n';
            }
        }
    }

    return 0;
}

#line 1 "test/src/datastructure/union_find/union_find_component_sum/abc183_f.test.cpp"
#define PROBLEM "https://atcoder.jp/contests/abc183/tasks/abc183_f"

#include <iostream>
#include <map>

#line 1 "library/datastructure/union_find/union_find_component_sum.hpp"



#line 1 "library/datastructure/union_find/union_find.hpp"



#include <algorithm>
#include <vector>

namespace suisen {
    struct UnionFind {
        UnionFind() = default;
        explicit UnionFind(int _n) : _n(_n), _dat(_n, -1) {}
        // Get the root of `x`. equivalent to `operator[](x)`
        int root(int x) {
            static std::vector<int> buf;
            while (_dat[x] >= 0) buf.push_back(x), x = _dat[x];
            while (buf.size()) _dat[buf.back()] = x, buf.pop_back();
            return x;
        }
        // Get the root of `x`. equivalent to `root(x)`
        int operator[](int x) {
            return root(x);
        }
        // Merge two vertices `x` and `y`.
        bool merge(int x, int y) {
            x = root(x), y = root(y);
            if (x == y) return false;
            if (_dat[x] > _dat[y]) std::swap(x, y);
            _dat[x] += _dat[y], _dat[y] = x;
            return true;
        }
        // Check if `x` and `y` belongs to the same connected component.
        bool same(int x, int y) {
            return root(x) == root(y);
        }
        // Get the size of connected component to which `x` belongs.
        int size(int x) {
            return -_dat[root(x)];
        }
        // Get all of connected components.
        std::vector<std::vector<int>> groups() {
            std::vector<std::vector<int>> res(_n);
            for (int i = 0; i < _n; ++i) res[root(i)].push_back(i);
            res.erase(std::remove_if(res.begin(), res.end(), [](const auto& g) { return g.empty(); }), res.end());
            return res;
        }
    protected:
        int _n;
        std::vector<int> _dat;
    };
} // namespace suisen


#line 5 "library/datastructure/union_find/union_find_component_sum.hpp"

namespace suisen {
    template <typename T, void(*merge_data)(T&, T)>
    struct UnionFindComponentSum : UnionFind {
        UnionFindComponentSum() : UnionFindComponentSum(0) {}
        explicit UnionFindComponentSum(int n, const T &init_value = T{}) : UnionFindComponentSum(std::vector<T>(n, init_value)) {}
        explicit UnionFindComponentSum(const std::vector<T> &init_values) : UnionFind(init_values.size()), _sum(init_values) {}

        bool merge(int x, int y) {
            x = root(x), y = root(y);
            bool res = UnionFind::merge(x, y);
            if (res) {
                if (root(x) == y) std::swap(x, y);
                merge_data(_sum[x], std::move(_sum[y]));
            }
            return res;
        }
        const T& sum(int x) {
            return _sum[root(x)];
        }
    private:
        std::vector<T> _sum;
    };
} // namespace suisen



#line 7 "test/src/datastructure/union_find/union_find_component_sum/abc183_f.test.cpp"

using sum_type = std::map<int, int>;

void merge(sum_type &par_data, sum_type ch_data) {
    for (auto &[k, v] : ch_data) par_data[k] += v;
}

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int n, q;
    std::cin >> n >> q;

    std::vector<int> c(n);
    for (int i = 0; i < n; ++i) std::cin >> c[i];

    std::vector<sum_type> init_data(n);
    for (int i = 0; i < n; ++i) {
        ++init_data[i][c[i]];
    }
    suisen::UnionFindComponentSum<sum_type, merge> uf(init_data);

    while (q --> 0) {
        int query_type;
        std::cin >> query_type;
        if (query_type == 1) {
            int a, b;
            std::cin >> a >> b;
            --a, --b;
            uf.merge(a, b);
        } else {
            int x, y;
            std::cin >> x >> y;
            --x;
            const auto& sum = uf.sum(x);
            if (auto it = sum.find(y); it == sum.end()) {
                std::cout << 0 << '\n';
            } else {
                std::cout << it->second << '\n';
            }
        }
    }

    return 0;
}