Test/String/TrieTree.test.cpp

• View this file on GitHub
• Last update: 2024-12-27 17:07:26+09:00
• Problem: https://yukicoder.me/problems/no/430

Depends on

• Library/String/TrieTree.hpp

Code

#define PROBLEM "https://yukicoder.me/problems/no/430"

#include "./../../Library/String/TrieTree.hpp"


#include <iostream>


using ll = long long;
using std::cin;
using std::cout;
constexpr char endl = '\n';

signed main() {
  std::string str;
  cin >> str;
  ll n;
  cin >> n;

  auto tree = mtd::TrieTree<bool, false>();
  for (int _ = 0; _ < n; ++_) {
    std::string s;
    cin >> s;
    std::vector<int> v(s.begin(), s.end());
    for (auto&& c : v) { c -= 'A'; }
    tree.emplace(v, true);
  }

  int sz = str.size();
  constexpr int ss = 10;
  ll ans = 0;
  for (int i = 0; i < sz; ++i) {
    std::vector<int> v;
    v.reserve(ss);
    for (int j = i; j < std::min(i + ss, sz); ++j) {
      v.emplace_back(str[j] - 'A');
    }
    tree.find(v, [&](bool) { ++ans; });
  }
  cout << ans << endl;
}

#line 1 "Test/String/TrieTree.test.cpp"
#define PROBLEM "https://yukicoder.me/problems/no/430"

#line 2 "Library/String/TrieTree.hpp"

#include <memory>

#include <vector>


namespace mtd {
  template <class Val = bool, Val ignore = false>
  class TrieTree {
    Val m_val;
    std::vector<std::unique_ptr<TrieTree>> m_next;
    static constexpr auto nullLambda = [](int) {};

    auto emplace(const std::vector<int>& vec, Val val, int it) {
      if (it == vec.size()) {
        m_val = val;
        return;
      }
      if (!m_next[vec[it]]) { m_next[vec[it]] = std::make_unique<TrieTree>(); }
      m_next[vec[it]]->emplace(vec, val, it + 1);
    }

    template <class Lambda>
    auto find(const std::vector<int>& vec, int it, const Lambda& lambda) const {
      if (m_val != ignore) { lambda(m_val); }
      if (it == vec.size()) { return m_val; }
      if (!m_next[vec[it]]) { return ignore; }
      return m_next[vec[it]]->find(vec, it + 1, lambda);
    }

  public:
    TrieTree() : TrieTree(ignore) {}
    TrieTree(Val val)
        : m_val(val), m_next(std::vector<std::unique_ptr<TrieTree>>(26)) {}

    auto emplace(const std::vector<int>& vec, Val val) {
      return emplace(vec, val, 0);
    }

    template <class Lambda = decltype(nullLambda)>
    auto find(const std::vector<int>& vec, const Lambda& lambda = nullLambda) {
      return find(vec, 0, lambda);
    }
  };
}  // namespace mtd

#line 4 "Test/String/TrieTree.test.cpp"

#include <iostream>


using ll = long long;
using std::cin;
using std::cout;
constexpr char endl = '\n';

signed main() {
  std::string str;
  cin >> str;
  ll n;
  cin >> n;

  auto tree = mtd::TrieTree<bool, false>();
  for (int _ = 0; _ < n; ++_) {
    std::string s;
    cin >> s;
    std::vector<int> v(s.begin(), s.end());
    for (auto&& c : v) { c -= 'A'; }
    tree.emplace(v, true);
  }

  int sz = str.size();
  constexpr int ss = 10;
  ll ans = 0;
  for (int i = 0; i < sz; ++i) {
    std::vector<int> v;
    v.reserve(ss);
    for (int j = i; j < std::min(i + ss, sz); ++j) {
      v.emplace_back(str[j] - 'A');
    }
    tree.find(v, [&](bool) { ++ans; });
  }
  cout << ans << endl;
}

Back to top page