AtCoder Grand Contest 047-B First Second (Trie木解法）

B - First Second

Trie木を使った実装を説明します。

実装については

algo-logic.info

を参考にしています。Trie木全体の実装などはこちらを読むと良いと思います。ここでは、この問題を解くために必要な部分だけを記述します。

問題概要

文字列がN個与えられる。文字列に対して次の「操作」を繰り返し行うことで新たに文字列を生成することができる。二つの文字列のペアであって片方からもう片方を「操作」によって作ることができるものは何通りあるか？

「操作」：文字列の先頭2文字のうち一方を削除する。

N <= 1e5 , Σ|S| <= 1e6

解法

1. 何をすれば良いか

①　長さMの文字列Sについて、生成できる文字列は、ある i (0 <= i <= M )について、[0 ; i ) の範囲から取ってきた1文字と、Sの接尾辞[i ; M )を連結させたものです。文字列Sをreverseしておきましょう。すると、Sの接頭辞およびそれ以降のうちから0文字または1文字を取ってきてつなげたものが生成されると言い換えることができます。

②　明らかに、自分より短い文字列から自分を生成することはできません。

まずは文字列たちを長さでソートしましょう。ソートの比較関数はグローバルに書くか、ソートの第3引数にラムダ記法を用いて

sort(a.begin(),a.end(),[](string& s,string& t){
  return s.size() < t.size();
});

と書くことで実行できます。

文字列Sから文字列Tが生成できるかを考えます。文字列Sから生成されうる文字列を効率的に全部列挙するにはどうすれば良いでしょうか？上で述べたように、操作によって生成される文字列はSの接頭辞 + alpha のような形をしています。接頭辞を、0文字、1文字、...と順番に生成して、接頭辞以降の文字列から取ってくることができる全ての文字を取ってくることで全部の生成文字列を列挙できます。

「接頭辞以降の文字列から取ってくることができる文字」はどうすればわかるでしょうか？アルファベットは26文字しかないので、26文字の小英文字全てについて、その文字を取ってくることができるかを試して良いです。つまり、ある文字cが、i文字目以降に出現するかどうかがわかれば良いです。これは、O(|S|)で前計算すれば簡単に求められます。(ABC-Cぐらいの水準？）ここでlogをつけると危ないんでしょうか...?

2.アルゴリズム

アルゴリズムは次の通りです。vector<string> A(N)に、文字列が長さの短い順に格納されています。

・string A[i]から生成される文字列を上に述べた方法で列挙する。ある生成文字列UがA[i-1],A[i-2],...,A[0]に存在するかをTrie木で判定。存在するなら答えに1加算。

・Trie木にA[i]を追加する。

上記をi = 0,1,...,N-1の順で繰り返す。

3.Trie木による文字列の存在判定

Trie木自体はいろいろと複雑な機能を搭載できますが、この問題を解くために必要な部分だけを記述します。まず、Trie木上の文字に対応するノードを次のように定義します。

struct Node{
    int c; //文字に対応した数字（'a'を引いた値）
    bool accept;  // ここで終わる文字列が存在（今回それは高々1つ）
    vector<int> nextid; // 自分から辺が張られる頂点番号(26文字に対応)
    Node(int c):c(c),accept(false){
      nextid.assign(26,-1);// -1 : 文字('a' + i)に対応する文字に辺がはられていない
    }
};

vector<Node> g; // グラフ

先に、文字列の挿入フェーズを記述しましょう。

void insert(string& s){
  int tmp = 0; // 現在の頂点番号
  for(int i = 0;i < s.size();i++){
    int& nid = g[tmp].nextid[s[i] - 'a']; //次に飛ぶべき頂点番号
    if(nid == -1){ // -1 : まだその頂点がグラフ上に存在しない
      nid = g.size(); // 新しく生成した頂点は、g[g.size()]になる
      g.emplace_back(s[i] - 'a');
    }
    tmp = nid; // 次の頂点に飛ぶ
  }
  g[tmp].accept = true; // 文字列の最後の文字はaccept（いわゆる受理状態）
}

文字列Sから生成される文字列たちが、Trie木上に存在するかを探索するパートをfunc(string s)として、次に記述します。

bool match(int tmp,int c){
    int nid = g[tmp].nextid[c]; // tmpから飛ぶ文字('a'+ c )に対応するノード番号
    return (nid != -1) && g[nid].accept; // nidで終わる文字列が存在するか
}
void func(string s){
    int sz = s.size();
    vector<int> last(26,-1);
    // 文字xが最後に登場する場所を前計算
    {
        for(int i = sz - 1;i >= 0;i--){
            if(last[s[i] - 'a'] == -1) last[s[i] - 'a'] - i;
        }
    }
    int tmp = 0; // 現在のノード番号
    for(int i = 0;i < sz;i++){
        // 接頭辞[0,i)に、[i,sz)のうち1文字を追加する
        for(int c = 0;c < 26;c++){
            if(j <= last[c] && match(tmp,c)){
                ans++;
            }
        }
        tmp = g[tmp].nextid[s[i] - 'a'];
　　　　　// [0,i+1)を接頭辞とする文字列が存在しないので、それ以上探索する必要がない
        if(tmp == -1) break; 
    }
    insert(s);
}

Trie木そのものの説明とか図とかは暇があったら描きたします、わかりにくいところなどありましたらご連絡お願いします。

提出： Submission #15813418 - AtCoder Grand Contest 047

コンテスト中に通せなかったのがとても心苦しいです。というかTrie木をちゃんと勉強していなかったので、脳裏には浮かんだけれどsetでぶん殴ろうとしてTLE...という悲しい結果でした。