:heavy_check_mark: test/Tree/trie/aoj-ALDS1_4.cpp

Depends on

Code

// competitive-verifier: PROBLEM https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ALDS1_4_C

#include "Tree/trie.hpp"
#include <bits/stdc++.h>

using namespace std;

int main()
{
    int n;
    cin >> n;

    Trie<char> trie;

    for (int i = 0; i < n; i++)
    {
        string S1, S2;
        cin >> S1 >> S2;
        if (S1 == "insert")
        {
            trie.insert(S2);
        }
        else
        {
            bool ans = false;
            auto find = [&](int id, bool is_end) -> void
            {
                if (is_end && 0 < trie.nodes[id].end_count)
                {
                    ans = true;
                }
            };

            trie.iterate_prefix(S2, find);
            cout << (ans ? "yes" : "no") << endl;
        }
    }

    return 0;
}
#line 1 "test/Tree/trie/aoj-ALDS1_4.cpp"
// competitive-verifier: PROBLEM https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ALDS1_4_C

#line 1 "Tree/trie.hpp"
#include <bits/stdc++.h>
using namespace std;

template <class T>
struct Trie
{
    struct Node
    {
        unordered_map<T, int> next_node_ids;
        int pass_count = 0;
        int end_count = 0;

        Node() {};
    };

    vector<Node> nodes;
    Trie() : nodes()
    {
        nodes.push_back(Node());
    };

    template <typename Container>
        requires same_as<typename Container::value_type, T>
    void insert(const Container &v)
    {
        int node_id = 0;
        nodes[node_id].pass_count++;
        for (const auto &c : v)
        {
            if (!nodes[node_id].next_node_ids.contains(c))
            {
                nodes[node_id].next_node_ids[c] = nodes.size();
                nodes.push_back(Node());
            }

            node_id = nodes[node_id].next_node_ids[c];
            nodes[node_id].pass_count++;
        }

        nodes[node_id].end_count++;
    }

    template <typename Container, typename Func>
        requires same_as<typename Container::value_type, T> && invocable<Func, int, bool>
    void iterate_prefix(const Container &v, Func f)
    {
        int node_id = 0;
        for (auto itr = v.begin(); itr != v.end(); itr++)
        {
            if (!nodes[node_id].next_node_ids.contains(*itr))
            {
                return;
            }

            node_id = nodes[node_id].next_node_ids[*itr];
            f(node_id, next(itr) == v.end());
        }
    }

    template <typename Container>
        requires same_as<typename Container::value_type, T>
    void erase_subtrie(const Container &prefix)
    {
        int node_id = 0;
        vector<int> ancestors;
        for (const auto &c : prefix)
        {
            ancestors.push_back(node_id);
            if (!nodes[node_id].next_node_ids.contains(c))
            {
                return;
            }

            node_id = nodes[node_id].next_node_ids[c];
        }

        int removed_pass_count = nodes[node_id].pass_count;
        nodes[node_id].pass_count = 0, nodes[node_id].end_count = 0;
        nodes[node_id].next_node_ids.clear();

        for (const auto &ancestor : ancestors)
        {
            nodes[ancestor].pass_count -= removed_pass_count;
        }
    }
};
#line 5 "test/Tree/trie/aoj-ALDS1_4.cpp"

using namespace std;

int main()
{
    int n;
    cin >> n;

    Trie<char> trie;

    for (int i = 0; i < n; i++)
    {
        string S1, S2;
        cin >> S1 >> S2;
        if (S1 == "insert")
        {
            trie.insert(S2);
        }
        else
        {
            bool ans = false;
            auto find = [&](int id, bool is_end) -> void
            {
                if (is_end && 0 < trie.nodes[id].end_count)
                {
                    ans = true;
                }
            };

            trie.iterate_prefix(S2, find);
            cout << (ans ? "yes" : "no") << endl;
        }
    }

    return 0;
}

Test cases

Env Name Status Elapsed Memory
g++ 01.in :heavy_check_mark: AC 5 ms 3 MB
g++ 02.in :heavy_check_mark: AC 4 ms 3 MB
g++ 03.in :heavy_check_mark: AC 4 ms 3 MB
g++ 04.in :heavy_check_mark: AC 4 ms 3 MB
g++ 05.in :heavy_check_mark: AC 4 ms 3 MB
g++ 06.in :heavy_check_mark: AC 4 ms 4 MB
g++ 07.in :heavy_check_mark: AC 11 ms 5 MB
g++ 08.in :heavy_check_mark: AC 17 ms 6 MB
g++ 09.in :heavy_check_mark: AC 1277 ms 60 MB
g++ 10.in :heavy_check_mark: AC 1450 ms 60 MB
Back to top page