【字典树】【KMP】【C++算法】3045统计前后缀下标对 II

template<class TData = char, int iTypeNum = 26, TData cBegin = 'a'>
class CTrieNode
{
public:
  CTrieNode* AddChar(TData ele, int& iMaxID)
  {
#ifdef _DEBUG
    if ((ele < cBegin) || (ele >= cBegin + iTypeNum))
    {
      return nullptr;
    }
#endif
    const int index = ele - cBegin;
    auto ptr = m_vPChilds[ele - cBegin];
    if (!ptr)
    {
      m_vPChilds[index] = new CTrieNode();
#ifdef _DEBUG
      m_vPChilds[index]->m_iID = ++iMaxID;
      m_childForDebug[ele] = m_vPChilds[index];
#endif
    }
    return m_vPChilds[index];
  }
  CTrieNode* GetChild(TData ele)const
  {
#ifdef _DEBUG
    if ((ele < cBegin) || (ele >= cBegin + iTypeNum))
    {
      return nullptr;
    }
#endif
    return m_vPChilds[ele - cBegin];
  }
protected:
#ifdef _DEBUG
  int m_iID = -1;
  std::unordered_map<TData, CTrieNode*> m_childForDebug;
#endif
public:
  int m_iLeafIndex = -1;
protected:
  CTrieNode* m_vPChilds[iTypeNum] = { nullptr };
};
template<class TData = char, int iTypeNum = 26, TData cBegin = 'a'>
class CTrie
{
public:
  int GetLeadCount()
  {
    return m_iLeafCount;
  }
  template<class IT>
  int Add(IT begin, IT end)
  {
    auto pNode = &m_root;
    for (; begin != end; ++begin)
    {
      pNode = pNode->AddChar(*begin, m_iMaxID);
    }
    if (-1 == pNode->m_iLeafIndex)
    {
      pNode->m_iLeafIndex = m_iLeafCount++;
    }
    return pNode->m_iLeafIndex;
  }
  template<class IT>
  CTrieNode<TData, iTypeNum, cBegin>* Search(IT begin, IT end)
  {
    auto ptr = &m_root;
    for (; begin != end; ++begin)
    {
      ptr = ptr->GetChild(begin);
      if (nullptr == ptr)
      {
        return nullptr;
      }
    }
    return ptr;
  }
  CTrieNode<TData, iTypeNum, cBegin> m_root;
protected:
  int m_iMaxID = 0;
  int m_iLeafCount = 0;
};
class KMP
{
public:
  virtual int Find(const string& s, const string& t)
  {
    CalLen(t);
    m_vSameLen.assign(s.length(), 0);
    for (int i1 = 0, j = 0; i1 < s.length(); )
    {
      for (; (j < t.length()) && (i1 + j < s.length()) && (s[i1 + j] == t[j]); j++);
      //i2 = i1 + j 此时s[i1,i2)和t[0,j)相等 s[i2]和t[j]不存在或相等
      m_vSameLen[i1] = j;
      //t[0,j)的结尾索引是j-1，所以最长公共前缀为m_vLen[j-1]，简写为y 则t[0,y)等于t[j-y,j)等于s[i2-y,i2)
      if (0 == j)
      {
        i1++;
        continue;
      }
      const int i2 = i1 + j;
      j = m_vLen[j - 1];
      i1 = i2 - j;//i2不变
    }
    for (int i = 0; i < m_vSameLen.size(); i++)
    {//多余代码是为了增加可测试性
      if (t.length() == m_vSameLen[i])
      {
        return i;
      }
    }
    return -1;
  }
  vector<int> m_vSameLen;//m_vSame[i]记录 s[i...]和t[0...]最长公共前缀，增加可调试性
  static vector<int> Next(const string& s)
  {
    const int len = s.length();
    vector<int> vNext(len, -1);
    for (int i = 1; i < len; i++)
    {
      int next = vNext[i - 1];
      while ((-1 != next) && (s[next + 1] != s[i]))
      {
        next = vNext[next];
      }
      vNext[i] = next + (s[next + 1] == s[i]);
    }
    return vNext;
  }
protected:
  void CalLen(const string& str)
  {
    m_vLen.resize(str.length());
    for (int i = 1; i < str.length(); i++)
    {
      int next = m_vLen[i - 1];
      while (str[next] != str[i])
      {
        if (0 == next)
        {
          break;
        }
        next = m_vLen[next-1];
      }
      m_vLen[i] = next + (str[next] == str[i]);
    }
  }
  int m_c;
  vector<int> m_vLen;//m_vLen[i] 表示t[0,i]的最长公共前后缀 
};
class Solution {
public:
  long long countPrefixSuffixPairs(vector<string>& words) {
    CTrie<> trie;
    unordered_map<int, int> mNoNum;
    long long llRet = 0;
    for (const auto& str : words)
    {     
      KMP kmp;
      kmp.Find(str, str);
      queue<int> indexs;
      for (int i = str.length()-1; i >= 0 ; i--)
      {
        if (kmp.m_vSameLen[i] == (str.length() - i))
        {
          indexs.emplace(str.length() - i);
        }
      }
      
      auto ptr = &trie.m_root;
      for (int i = 0; i < str.length(); i++)
      {
        ptr = ptr->GetChild(str[i]);
        if (nullptr == ptr)
        {
          break;
        }
        if ((-1 != ptr->m_iLeafIndex)&&indexs.size()&&( indexs.front()==i+1 ))
        {
          llRet += mNoNum[ptr->m_iLeafIndex];         
        }
        while (indexs.size() && (indexs.front() == i + 1))
        {
          indexs.pop();
        }
      }
      mNoNum[trie.Add(str.begin(), str.end())]++;
    }
    return llRet;
  }
};