fix bugs and modified aho-corasick

string/aho-corasick.hpp

@@ -2,15 +2,15 @@
 
 #include "trie.hpp"
 
-template <size_t X = 26, char margin = 'a'>
+template <size_t X = 26, char margin = 'a', bool heavy = false>
 struct AhoCorasick : Trie<X + 1, margin> {
   using TRIE = Trie<X + 1, margin>;
   using TRIE::next;
   using TRIE::st;
   using TRIE::TRIE;
   vector<int> cnt;
 
-  void build(int heavy = true) {
+  void build() {
     int n = st.size();
     cnt.resize(n);
     for (int i = 0; i < n; i++) {
@@ -55,15 +55,24 @@ struct AhoCorasick : Trie<X + 1, margin> {
     }
   }
 
-  vector<int> match(string s, int heavy = true) {
-    vector<int> res(heavy ? TRIE::size() : 1);
+  // heavy
+  // true  : 各パターン文字列に対してマッチした回数を計算
+  // false : 全てのパターン文字列にマッチした回数の総和
+  conditional_t<heavy, unordered_map<int, long long>, long long> match(
+      string s) {
+    unordered_map<int, int> pos_cnt;
     int pos = 0;
     for (auto &c : s) {
       pos = next(pos, c - margin);
-      if (heavy)
-        for (auto &x : st[pos].idxs) res[x]++;
-      else
-        res[0] += cnt[pos];
+      pos_cnt[pos]++;
+    }
+    conditional_t<heavy, unordered_map<int, long long>, long long> res{};
+    for (auto &[key, val] : pos_cnt) {
+      if constexpr (heavy) {
+        for (auto &x : st[key].idxs) res[x] += val;
+      } else {
+        res += 1LL * cnt[key] * val;
+      }
     }
     return res;
   }

verify/verify-yuki/yuki-0430-2.test.cpp

@@ -0,0 +1,27 @@
+#define PROBLEM "https://yukicoder.me/problems/no/430"
+//
+#include "../../template/template.hpp"
+//
+#include "../../string/aho-corasick.hpp"
+using namespace Nyaan;
+
+void q() {
+  ins(S);
+  ini(M);
+  AhoCorasick<26, 'A', true> ac;
+  rep(i, M) {
+    ins(C);
+    ac.add(C, i);
+  }
+  ac.build();
+  auto ans = ac.match(S);
+  ll bns = 0;
+  each2(key, val, ans) bns += val;
+  out(bns);
+}
+
+void Nyaan::solve() {
+  int t = 1;
+  // in(t);
+  while (t--) q();
+}

verify/verify-yuki/yuki-0430.test.cpp

@@ -0,0 +1,24 @@
+#define PROBLEM "https://yukicoder.me/problems/no/430"
+//
+#include "../../template/template.hpp"
+//
+#include "../../string/aho-corasick.hpp"
+using namespace Nyaan;
+
+void q() {
+  ins(S);
+  ini(M);
+  AhoCorasick<26, 'A', false> ac;
+  rep(i, M) {
+    ins(C);
+    ac.add(C, 0);
+  }
+  ac.build();
+  out(ac.match(S));
+}
+
+void Nyaan::solve() {
+  int t = 1;
+  // in(t);
+  while (t--) q();
+}

verify/verify-yuki/yuki-1269.test.cpp

@@ -23,7 +23,7 @@ void Nyaan::solve() {
 
   AhoCorasick<10, '0'> aho;
   rep(i, sz(ss)) aho.add(ss[i], i);
-  aho.build(false);
+  aho.build();
 
   vm dp(aho.size());
   dp[0] = 1;