Suffix Arrayを使ってみた

全く解き方が分からなかったMountain HolidaysのEditorialが出たので解いてみました．

Mountain Holidays

整数の配列height(1 <= |height| <= 100000)が与えられる．このとき，連続部分列の数を数えよ．

次の条件のいずれかを満たすとき(i1, j1)と(i2, j2)の2つの連続部分列は異なる．
1. j1 – i1 != j2 – i2
2. height[i1+k] – height[i1+k-1] != height[i2+k] – height[i2+k–1] を満たすkが (1, 2, ... j1-i1)の範囲に存在．

この問題は異なる部分文字列の個数を数える問題に帰着できます．height[i] - height[i-1]をその整数に対応する文字と考えます．
しかし，この問題では|height|が大きいため，ハッシュなどを使用して高速な探索を行ったとしても部分文字列の個数を探すのにO(N^2)かかってしまいTLEします．

そこで，heightのsuffixを辞書順に並べるSuffix Arrayを構築します．Suffix Arrayを用いることでi番目のsuffixとi+1番目のsuffixのLCP(Longest Common Prefix)を高速に求めることができ，（全ての部分文字列の個数） - （i番目のsuffixとi+1番目のsuffixのLCPの総和）で異なる部分文字列の個数を出すことが出来ます．

Suffix Arrayの構築は普通にやると，ソートにO(NlogN)，1回の比較にO(N)かかるため，O(N^2logN)となり間に合いません．
そこで，2^iの長さごとに辞書順を求めておき，2^(i+1)の長さの辞書順を求める際に2^iの長さの辞書順を用いることでO(N(logN)^2)でSuffix Arrayを得ることが出来ます．Suffix Arrayを得ることが出来れば1つのLCPをO(logN)で求められるため，LCPの総和をO(NlogN)で求めることが出来，TLEせずに実行することが出来ます．

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <vector>
#include <algorithm>
#include <string>
#include <map>
#include <set>
#include <queue>
#include <stack>
#include <climits>
#include <sstream>
#include <functional>
#include <complex>

using namespace std;

#define len(array)  (sizeof (array) / sizeof *(array))
#define rep(i, s, e) for(int i = s;i < e;i++)
#define Rep(i, e) for(int i = 0;i < e;i++)
#define rrep(i, e, s) for(int i = e;s <= i;i--)
#define Rrep(i, e) for(int i = e;0 <= i;i--)
#define mrep(i, e, t1, t2) for(map<t1, t2>::iterator i = e.begin(); i != e.end(); i++)
#define vrange(v) v.begin(), v.end()
#define vrrange(v) v.rbegin(), v.rend()
#define vsort(v) sort(vrange(v))
#define vrsort(v) sort(vrrange(v))
#define arange(a) a, a + len(a)
#define asort(a) sort(arange(a))
#define arsort(a, t) sort(arange(a), greater<t>())
#define afill(a, v) fill(arange(a), v)
#define afill2(a, v, t) fill((t *)a, (t *)(a + len(a)), v)
#define fmax(a, b) (a < b? b : a)
#define fmin(a, b) (a > b? b : a)
#define fabs(a) (a < 0? -a : a)
#define pb push_back
#define rg(i, s, t) s <= i && i < t
//#define X real()
//#define Y imag()
//typedef unsigned int ui;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int, int> P;
//typedef complex<double> p;
const int INF = (int)2e9;
const int MOD = (int)1e9 + 9;
const double EPS = 1e-10;
//const int dx[] = {1, -1, 0, 0, 1, -1, -1, 1};
//const int dy[] = {0, 0, 1, -1, -1, -1, 1, 1};
//const int weight[] = {0,1,10,100,1000,10000,100000,1000000,10000000};
//priority_queue< int, vector<int>, greater<int> > q;

//Suffix Array
//http://www.stanford.edu/class/cs97si/suffix-array.pdf
//http://discuss.codechef.com/questions/18105/mou1h-editorial

#define MAX_N 100005
#define MAXLG 18
#define MIN_VALUE -INF
struct entry {
  int nr[2], p;

  bool operator <(const struct entry &e) const{
    return nr[0] == e.nr[0] ? (nr[1] == e.nr[1]? p < e.p : nr[1] < e.nr[1]) : (nr[0] < e.nr[0]);
  }
  bool operator >(const struct entry &e) const{
    return nr[0] == e.nr[0] ? (nr[1] == e.nr[1]? p > e.p : nr[1] > e.nr[1]) : (nr[0] > e.nr[0]);
  }
} L[MAX_N];
int SA[MAXLG][MAX_N], N, suffix[MAX_N], LCP[MAX_N];
int A[MAX_N], D[MAX_N];
//SA[k][i] : the order of the subsequence of A of length 2^k starting at position i.
//return the last row index of matrix SA
// O(N (logN)^2)
int makeSuffixArray(int n) {
  // cin >> A;
  int stp = 1, cnt = 1;
  N = n;
  Rep(i, N) SA[0][i] = D[i];
  while(cnt < N){
    Rep(i, N){
      L[i].nr[0] = SA[stp - 1][i];
      L[i].nr[1] = (i + cnt < N)? SA[stp - 1][i + cnt] : MIN_VALUE;
      L[i].p = i;
    }
    sort(L, L + N);
    Rep(i, N){
      SA[stp][L[i].p] = (i > 0 && L[i].nr[0] == L[i - 1].nr[0] && L[i].nr[1] == L[i - 1].nr[1])? SA[stp][L[i - 1].p] : i;
    }
    stp++;
    cnt *= 2;
  }
  return stp-1;
}

//Computing the longest common prefix (LCP)
// O(log N)
int lcp(int x, int y, int stp)
{
  int ret = 0;
  if (x == y) return N - x;
  for (int k = stp; k >= 0 && x < N && y < N; k--){
    if (SA[k][x] == SA[k][y]){
      x += (1 << k);
      y += (1 << k);
      ret += (1 << k);
    }
  }
  return ret;
}

//結果はsuffixに格納．
int getSuffixArray(int n){
  int idx = makeSuffixArray(n);
  Rep(i, N){
    suffix[SA[idx][i]] = i;
  }
  return idx;
}

//結果はLCPに格納
void getLCPArray(int stp){
  LCP[0] = 0;
  rep(i, 1, N){
    LCP[i] = lcp(suffix[i-1], suffix[i], stp);
  }
}

void doIt(){
  int t, n;
  ll ans, sum;
  scanf("%d", &t);
  while(t--){
    scanf("%d", &n);
    Rep(i, n) scanf("%d", &A[i]);
    if(n == 1){
      printf("0\n");
      continue;
    }
    else if(n == 2){
      printf("1\n");
      continue;
    }
    ans = n;
    ans *= n-1;
    ans /= 2;
    ans %= MOD;
    n--;
    rep(i, 0, n) D[i] = A[i+1] - A[i];
    int idx = getSuffixArray(n);
    getLCPArray(idx);
    // printf("idx = %d\n", idx);
    // printf("idx2 = %f\n", ceil(log(N)/log(2)));
    Rep(i, n){
      ans -= LCP[i];
      // printf("LCP[%d] = %d\n", i, LCP[i]);
      if(ans < 0) ans += MOD;
      ans %= MOD;
    }
    printf("%lld\n", ans);
  }


}

int main() {
  doIt();
  return 0;
}

この問題，全く通せなくてすごい苦労しました．Editorial通りに実装しているし，他の人のコードを見ても同じようにやっているのになんで通らないんだ？と数時間格闘し，最後は自力は諦めて通っている人のコードと差分を取りながらデバッグ．
その結果，MODに使う数字がずれていたと判明（！！？）．多くの問題ではMODは10^9+7なのですが，この問題は10^9+9というしょうもないオチでした．これはひどい．脱力・・・．

参考
http://discuss.codechef.com/questions/18105/mou1h-editorial
http://www.stanford.edu/class/cs97si/suffix-array.pdf