[Codechef] PALIN3 - 3 Palindromes

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : PALIN3 - 3 Palindromes
Source            : Codechef
Category          : Data Structure
Algorithm         : Palindromic Tree
Verdict           : Accepted

#include <bits/stdc++.h>  
  
using namespace std;  
  
#define ll              long long int  
  
static const int maxn = 1e6 + 5;  
  
struct Node  
{  
      int Next[10];  
      int len;  
      int sufflink;  
      ll num;  
};  
  
string s;  
int len, suff, num;  
Node Tree[maxn * 4];  
  
int cumsum[maxn];  
  
int isDivisible(int l, int r)  
{  
      int sze = r - l + 1;  
      if (sze > 1 && s[l] == '0') return 0;  
      int sum = cumsum[r];  
      if (l-1 >= 0) sum -= cumsum[l-1];  
      return (sum % 3 == 0);  
}  
  
void init()  
{  
      num = 2; suff = 2;  
      Tree[1].len = -1; Tree[1].sufflink = 1;  
      Tree[2].len = 0; Tree[2].sufflink = 1;  
}  
  
bool addLetter(int pos)  
{  
      int cur = suff, curlen = 0;  
      int let = s[pos] - '0';  
      while (true)  
      {  
            curlen = Tree[cur].len;  
            if (pos - 1 - curlen >= 0 && s[pos - 1 - curlen] == s[pos]) break;  
            cur = Tree[cur].sufflink;  
      }  
      if (Tree[cur].Next[let])  
      {  
            suff = Tree[cur].Next[let];  
            return false;  
      }  
      num++;  
      suff = num;  
      Tree[num].len = Tree[cur].len + 2;  
      Tree[cur].Next[let] = num;  
      int add = isDivisible(pos - Tree[num].len + 1, pos);  
      if (Tree[num].len == 1)  
      {  
            Tree[num].sufflink = 2;  
            Tree[num].num = add;  
            return true;  
      }  
      while (true)  
      {  
            cur = Tree[cur].sufflink;  
            curlen = Tree[cur].len;  
            if (pos - 1 - curlen >= 0 && s[pos - 1 - curlen] == s[pos])  
            {  
                  Tree[num].sufflink = Tree[cur].Next[let];  
                  break;  
            }  
      }  
      Tree[num].num = add + Tree[ Tree[num].sufflink ].num;  
      return true;  
}  
  
signed main()  
{  
      ios_base::sync_with_stdio(false);  
      cin.tie(nullptr);  
      cout.tie(nullptr);  
      cout << fixed << setprecision(3);  
  
      #ifndef ONLINE_JUDGE  
            freopen("in.txt", "r", stdin);  
      #endif // ONLINE_JUDGE  
  
      cin >> s;  
      len = s.size();  
  
      cumsum[0] = s[0] - '0';  
      for (int i = 1; i < len; i++) cumsum[i] = (s[i] - '0') + cumsum[i-1];  
  
      init();  
      ll ans = 0;  
      for (int i = 0; i < len; i++)  
      {  
            addLetter(i);  
            ans += Tree[suff].num;  
      }  
      cout << ans << '\n';  
  
}