test/yosupo/associative-array.test.cpp

• View this file on GitHub
• Last update: 2026-10-13 06:02:01+09:00
• Problem: https://judge.yosupo.jp/problem/associative_array

Depends on

• 値の個数を数える多重集合 (Counter) (src/data_structures/counter.hpp)

Code

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

#include <iostream>

#include "../../src/data_structures/counter.hpp"

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    int q;
    std::cin >> q;
    // 値 v を「キー k の個数」として持つ (0 なら取り除かれる)
    Counter<long long> a;
    while (q--) {
        int t;
        long long k;
        std::cin >> t >> k;
        if (t == 0) {
            long long v;
            std::cin >> v;
            a[k] = v;
        } else {
            std::cout << a[k] << '\n';
        }
    }
}

#line 1 "test/yosupo/associative-array.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/associative_array"

#include <iostream>

#line 2 "src/data_structures/counter.hpp"

/**
 * @file counter.hpp
 * @brief 値の個数を数える多重集合 (Counter)
 */

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <initializer_list>
#include <iterator>
#include <map>
#include <type_traits>
#include <utility>
#include <vector>

template <typename T>
struct Counter {
    using key_type = T;
    using mapped_type = long long;
    using value_type = std::pair<key_type, mapped_type>;
    std::map<key_type, mapped_type> data;

    void inc(const key_type& key, const mapped_type& t) {
        data[key] += t;
        auto itr = data.find(key);
        assert(itr->second >= 0);
        if (itr->second == 0) {
            data.erase(itr);
        }
    }
    Counter() = default;
    explicit Counter(const key_type& x) : data{{x, 1}} {}
    /**
     * @brief 範囲の要素を数える
     * Counter 自身と key_type (std::string など範囲でもあるもの) はこちらで受け取らない
     */
    template <class C, std::enable_if_t<!std::is_same_v<std::decay_t<C>, Counter> && !std::is_convertible_v<C, key_type>, std::nullptr_t> = nullptr,
              class = decltype(std::begin(std::declval<C&>()), std::end(std::declval<C&>()))>
    explicit Counter(C&& v) {
        for (const key_type& a : v) {
            data[a]++;
        }
    }
    Counter(std::initializer_list<value_type> init) {
        for (const auto& p : init) {
            inc(p.first, p.second);
        }
    }
    Counter& operator+=(const Counter& o) {
        for (const auto& [i, c] : o.data) {
            data[i] += c;
        }
        return *this;
    }
    friend Counter operator+(const Counter& a, const Counter& b) {
        Counter res(a);
        return res += b;
    }
    Counter& operator-=(const Counter& o) {
        // 自分自身を引くと、走査中の o.data の要素を inc が消してしまう
        if (this == &o) {
            data.clear();
            return *this;
        }
        for (const auto& [i, c] : o.data) {
            inc(i, -c);
        }
        return *this;
    }
    friend Counter operator-(const Counter& a, const Counter& b) {
        Counter res(a);
        return res -= b;
    }
    auto begin() const { return data.begin(); }
    auto end() const { return data.end(); }
    std::size_t size() const {
        return data.size();
    }
    mapped_type get(const key_type& key) const {
        auto itr = data.find(key);
        if (itr == data.end()) return 0;
        return itr->second;
    }
    void set(const key_type& key, const mapped_type& t) {
        assert(t >= 0);
        if (t > 0) {
            data[key] = t;
        } else {
            auto itr = data.find(key);
            if (itr != data.end()) {
                data.erase(itr);
            }
        }
    }
    std::vector<value_type> most_common() const {
        std::vector<value_type> ret;
        ret.reserve(data.size());
        for (const auto& [i, c] : data) {
            ret.push_back({i, c});
        }
        std::sort(ret.begin(), ret.end(), [](const value_type& a, const value_type& b) { return a.second > b.second or (a.second == b.second and a.first < b.first); });
        return ret;
    }
    class CounterReference {
        Counter& base;
        key_type key;

       public:
        CounterReference(Counter& base, key_type key) : base(base), key(key) {}
        CounterReference& operator=(const mapped_type& x) {
            base.set(key, x);
            return *this;
        }
        CounterReference& operator+=(const mapped_type& x) {
            base.inc(key, x);
            return *this;
        }
        CounterReference& operator-=(const mapped_type& x) {
            base.inc(key, -x);
            return *this;
        }
        operator mapped_type() const {
            return base.get(key);
        }
    };
    mapped_type operator[](const key_type& key) const {
        return get(key);
    }
    CounterReference operator[](const key_type& key) {
        return CounterReference(*this, key);
    }
};
#line 6 "test/yosupo/associative-array.test.cpp"

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    int q;
    std::cin >> q;
    // 値 v を「キー k の個数」として持つ (0 なら取り除かれる)
    Counter<long long> a;
    while (q--) {
        int t;
        long long k;
        std::cin >> t >> k;
        if (t == 0) {
            long long v;
            std::cin >> v;
            a[k] = v;
        } else {
            std::cout << a[k] << '\n';
        }
    }
}

Back to top page