#608147 (C++17) No.430 文字列検索

• 問題: No.430 文字列検索
• ユーザー: beet
• 提出日時: 2021-01-22 18:34:44
• 言語: C++17 (gcc 13.3.0 + boost 1.87.0)
• 結果: RE
• 実行時間: -
• コード長: 3,562 bytes
• コンパイル時間: 1,991 ms
• コンパイル使用メモリ: 211,100 KB
• 最終ジャッジ日時: 2025-01-18 03:28:14
• ジャッジサーバーID （参考情報）: judge2 / judge4

ファイルパターン	結果
sample	AC * 3 RE * 1
other	AC * 10 RE * 4

ソースコード

// verification-helper: PROBLEM https://yukicoder.me/problems/1013

#include <bits/stdc++.h>
using namespace std;

#define call_from_test
template<typename T=int>
vector<T> read(size_t n){
  vector<T> ts(n);
  for(size_t i=0;i<n;i++) cin>>ts[i];
  return ts;
}

template<size_t X>
struct Trie{
  struct Node{
    array<int, X> nxt;
    vector<int> idxs;
    int idx;
    Node():idx(-1){fill(nxt.begin(),nxt.end(),-1);}
  };

  using F = function<int(char)>;
  vector<Node> vs;
  F conv;

  Trie(F conv):conv(conv){vs.emplace_back();}
  Trie(char start):Trie([=](char a){return a-start;}){}

  inline int &next(int i,int j){
    return vs[i].nxt[j];
  }

  void add(const string &s,int x){
    int pos=0;
    for(char c:s){
      int k=conv(c);
      if(~next(pos,k)){
        pos=next(pos,k);
        continue;
      }
      int npos=vs.size();
      next(pos,k)=npos;
      vs.emplace_back();
      pos=npos;
    }
    vs[pos].idx=x;
    vs[pos].idxs.emplace_back(x);
  }

  int find(const string &s){
    int pos=0;
    for(char c:s){
      int k=conv(c);
      if(next(pos,k)<0) return -1;
      pos=next(pos,k);
    }
    return pos;
  }

  int move(int pos,char c){
    assert(pos<(int)vs.size());
    return pos<0?-1:next(pos,conv(c));
  }

  int size(){return vs.size();}

  int idx(int pos){
    return pos<0?-1:vs[pos].idx;
  }

  vector<int> idxs(int pos){
    return pos<0?vector<int>():vs[pos].idxs;
  }
};

template<size_t X, bool heavy>
struct AhoCorasick : Trie<X+1>{
  using super = Trie<X+1>;
  using super::super, super::next, super::size;
  using super::vs, super::conv;

  vector<int> cnt,order;

  // O(\sigma \sum |T_i|)
  void build(){
    int n=vs.size();
    cnt.resize(n);
    for(int i=0;i<n;i++){
      if(heavy) sort(vs[i].idxs.begin(),vs[i].idxs.end());
      cnt[i]=vs[i].idxs.size();
    }

    queue<int> que;
    for(int i=0;i<(int)X;i++){
      if(~next(0,i)){
        next(next(0,i),X)=0;
        que.emplace(next(0,i));
      }else{
        next(0,i)=0;
      }
    }

    while(!que.empty()){
      order.emplace_back(que.front());
      auto &x=vs[que.front()];
      int fail=x.nxt[X];

      cnt[que.front()]+=cnt[fail];
      que.pop();

      for(int i=0;i<(int)X;i++){
        int &nx=x.nxt[i];
        if(nx<0){
          nx=next(fail,i);
          continue;
        }
        que.emplace(nx);
        next(nx,X)=next(fail,i);
        if(heavy){
          auto &idx=vs[nx].idxs;
          auto &idy=vs[next(fail,i)].idxs;
          vector<int> idz;
          set_union(idx.begin(),idx.end(),
                    idy.begin(),idy.end(),
                    back_inserter(idz));
          idx=idz;
        }
      }
    }
  }

  int count(int pos){
    return cnt[pos];
  }

  // O(|S|)
  long long match(string s){
    long long res=0;
    int pos=0;
    for(auto &c:s){
      pos=next(pos,conv(c));
      res+=count(pos);
    }
    return res;
  }

  // O(|S| + \sum |T_i|)
  vector<int> frequency(string s){
    vector<int> res(size(),0);
    int pos=0;
    for(auto &c:s){
      pos=next(pos,conv(c));
      res[pos]++;
    }
    for(int i=size()-1;i;i--)
      res[vs[order[i]].nxt[X]]+=res[order[i]];
    return res;
  }
};

#undef call_from_test

signed main(){
  cin.tie(0);
  ios::sync_with_stdio(0);

  AhoCorasick<26, false> aho('A');
  string s;
  cin>>s;
  int m;
  cin>>m;
  auto cs=read<string>(m);
  for(auto c:cs) aho.add(c,0);
  aho.build();
  cout<<aho.match(s)<<endl;

  auto res=aho.frequency(s);
  long long cnt=0;
  for(auto c:cs) cnt+=res[aho.find(c)];
  assert(cnt==aho.match(s));
  return 0;
}