From 92e55cbea52ae147e6aa96ef2894db49702d28b2 Mon Sep 17 00:00:00 2001
From: NyaanNyaan <suteakadapyon3@gmail.com>
Date: Wed, 7 Aug 2024 23:00:26 +0900
Subject: [PATCH] fix bugs and modified aho-corasick

---
 string/aho-corasick.hpp                 | 25 +++++++++++++++--------
 verify/verify-yuki/yuki-0430-2.test.cpp | 27 +++++++++++++++++++++++++
 verify/verify-yuki/yuki-0430.test.cpp   | 24 ++++++++++++++++++++++
 verify/verify-yuki/yuki-1269.test.cpp   |  2 +-
 4 files changed, 69 insertions(+), 9 deletions(-)
 create mode 100644 verify/verify-yuki/yuki-0430-2.test.cpp
 create mode 100644 verify/verify-yuki/yuki-0430.test.cpp

diff --git a/string/aho-corasick.hpp b/string/aho-corasick.hpp
index b6ea01907..cdbbb7697 100644
--- a/string/aho-corasick.hpp
+++ b/string/aho-corasick.hpp
@@ -2,7 +2,7 @@
 
 #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;
@@ -10,7 +10,7 @@ struct AhoCorasick : Trie<X + 1, margin> {
   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;
   }
diff --git a/verify/verify-yuki/yuki-0430-2.test.cpp b/verify/verify-yuki/yuki-0430-2.test.cpp
new file mode 100644
index 000000000..ed0adefef
--- /dev/null
+++ b/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();
+}
diff --git a/verify/verify-yuki/yuki-0430.test.cpp b/verify/verify-yuki/yuki-0430.test.cpp
new file mode 100644
index 000000000..38d7be568
--- /dev/null
+++ b/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();
+}
diff --git a/verify/verify-yuki/yuki-1269.test.cpp b/verify/verify-yuki/yuki-1269.test.cpp
index ac8cd306b..da6aee49d 100644
--- a/verify/verify-yuki/yuki-1269.test.cpp
+++ b/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;