Longest Common Extension / LCE using Segment Tree

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

// Structure to represent a query of form (L,R)
struct Query
{
    int L, R;
};

// Structure to store information of a suffix
struct suffix
{
    int index;  // To store the original index
    int rank[2]; // To store ranks and the next rank pair
};

// A utility function to get the minimum of two numbers
int minVal(int x, int y)
{
    return (x < y) ? x : y;
}

// A utility function to get the middle index from
// corner indexes.
int getMid(int s, int e)
{
    return s + (e - s) / 2;
}

/*  A recursive function to get the minimum value
    in a given range of array indexes. The following
    are parameters for this function.

    st    --> Pointer to the segment tree
    index --> Index of the current node in the segment
              tree. Initially, 0 is passed as the root
              is always at index 0
    ss & se  --> Starting and ending indexes of the
                  segment represented by the current
                  node, i.e., st[index]
    qs & qe  --> Starting and ending indexes of the query
                 range */
int RMQUtil(int *st, int ss, int se, int qs, int qe,
            int index)
{
    // If the segment of this node is a part of the given range,
    // then return the min of the segment
    if (qs <= ss && qe >= se)
        return st[index];

    // If the segment of this node is outside the given range
    if (se < qs || ss > qe)
        return INT_MAX;

    // If a part of this segment overlaps with the given range
    int mid = getMid(ss, se);
    return minVal(RMQUtil(st, ss, mid, qs, qe, 2 * index + 1),
                  RMQUtil(st, mid + 1, se, qs, qe, 2 * index + 2));
}

// Return the minimum of elements in the range from index qs
// (query start) to qe (query end). It mainly uses RMQUtil()
int RMQ(int *st, int n, int qs, int qe)
{
    // Check for erroneous input values
    if (qs < 0 || qe > n - 1 || qs > qe)
    {
        cout << "Invalid Input";
        return -1;
    }

    return RMQUtil(st, 0, n - 1, qs, qe, 0);
}

// A recursive function that constructs Segment Tree
// for array[ss..se]. si is the index of the current node in
// the segment tree st
int constructSTUtil(int arr[], int ss, int se, int *st,
                    int si)
{
    // If there is one element in the array, store it in
    // the current node of the segment tree and return
    if (ss == se)
    {
        st[si] = arr[ss];
        return arr[ss];
    }

    // If there are more than one element, then recur
    // for left and right subtrees and store the minimum
    // of two values in this node
    int mid = getMid(ss, se);
    st[si] = minVal(constructSTUtil(arr, ss, mid, st, si * 2 + 1),
                   constructSTUtil(arr, mid + 1, se, st, si * 2 + 2));
    return st[si];
}

/* Function to construct a segment tree from the given array.
   This function allocates memory for the segment tree and
   calls constructSTUtil() to fill the allocated memory */
int *constructST(int arr[], int n)
{
    // Allocate memory for the segment tree

    // Height of the segment tree
    int x = (int)(ceil(log2(n)));

    // Maximum size of the segment tree
    int max_size = 2 * (int)pow(2, x) - 1;

    int *st = new int[max_size];

    // Fill the allocated memory st
    constructSTUtil(arr, 0, n - 1, st, 0);

    // Return the constructed segment tree
    return st;
}

// A comparison function used by sort() to compare
// two suffixes. Compares two pairs, returns 1 if
// the first pair is smaller
int cmp(struct suffix a, struct suffix b)
{
    return (a.rank[0] == b.rank[0]) ?
               (a.rank[1] < b.rank[1] ? 1 : 0) :
               (a.rank[0] < b.rank[0] ? 1 : 0);
}

// This is the main function that takes a string
// 'txt' of size n as an argument, builds and returns
// the suffix array for the given string
vector<int> buildSuffixArray(string txt, int n)
{
    // A structure to store suffixes and their indexes
    struct suffix suffixes[n];

    // Store suffixes and their indexes in an array
    // of structures. The structure is needed to sort
    // the suffixes alphabetically and maintain their
    // old indexes while sorting
    for (int i = 0; i < n; i++)
    {
        suffixes[i].index = i;
        suffixes[i].rank[0] = txt[i] - 'a';
        suffixes[i].rank[1] = ((i + 1) < n) ?
                                 (txt[i + 1] - 'a') : -1;
    }

    // Sort the suffixes using the comparison function
    // defined above.
    sort(suffixes, suffixes + n, cmp);

    // At this point, all suffixes are sorted according to the first
    // 2 characters.  Let us sort suffixes according to the first 4
    // characters, then the first 8 and so on
    int ind[n]; // This array is needed to get the index
                // in suffixes[]
                // from the original index.  This mapping is needed to get
                // the next suffix.
    for (int k = 4; k < 2 * n; k = k * 2)
    {
        // Assigning rank and index values to the first suffix
        int rank = 0;
        int prev_rank = suffixes[0].rank[0];
        suffixes[0].rank[0] = rank;
        ind[suffixes[0].index] = 0;

        // Assigning rank to suffixes
        for (int i = 1; i < n; i++)
        {
            // If the first rank and next ranks are the same as
            // that of the previous suffix in the array, assign
            // the same new rank to this suffix
            if (suffixes[i].rank[0] == prev_rank &&
                suffixes[i].rank[1] == suffixes[i - 1].rank[1])
            {
                prev_rank = suffixes[i].rank[0];
                suffixes[i].rank[0] = rank;
            }
            else // Otherwise increment rank and assign
            {
                prev_rank = suffixes[i].rank[0];
                suffixes[i].rank[0] = ++rank;
            }
            ind[suffixes[i].index] = i;
        }

        // Assign the next rank to every suffix
        for (int i = 0; i < n; i++)
        {
            int nextindex = suffixes[i].index + k / 2;
            suffixes[i].rank[1] = (nextindex < n) ?
                                     suffixes[ind[nextindex]].rank[0] : -1;
        }

        // Sort the suffixes according to the first k characters
        sort(suffixes, suffixes + n, cmp);
    }

    // Store indexes of all sorted suffixes in the suffix array
    vector<int> suffixArr;
    for (int i = 0; i < n; i++)
        suffixArr.push_back(suffixes[i].index);

    // Return the suffix array
    return suffixArr;
}

/* To construct and return LCP */
vector<int> kasai(string txt, vector<int> suffixArr,
                  vector<int> &invSuff)
{
    int n = suffixArr.size();

    // To store the LCP array
    vector<int> lcp(n, 0);

    // Fill values in invSuff[]
    for (int i = 0; i < n; i++)
        invSuff[suffixArr[i]] = i;

    // Initialize the length of the previous LCP
    int k = 0;

    // Process all suffixes one by one starting from
    // the first suffix in txt[]
    for (int i = 0; i < n; i++)
    {
        /* If the current suffix is at n-1, then we don?t
           have the next substring to consider. So the LCP is not
           defined for this substring, we put zero. */
        if (invSuff[i] == n - 1)
        {
            k = 0;
            continue;
        }

        /* j contains the index of the next substring to
           be considered  to compare with the present
           substring, i.e., the next string in the suffix array */
        int j = suffixArr[invSuff[i] + 1];

        // Directly start matching from the k'th index as
        // at least k-1 characters will match
        while (i + k < n && j + k < n && txt[i + k] == txt[j + k])
            k++;

        lcp[invSuff[i]] = k; // LCP for the present suffix.

        // Deleting the starting character from the string.
        if (k > 0)
            k--;
    }

    // return the constructed LCP array
    return lcp;
}

// A utility function to find the longest common extension
// from index - L and index - R
int LCE(int *st, vector<int> lcp, vector<int> invSuff,
        int n, int L, int R)
{
    // Handle the corner case
    if (L == R)
        return (n - L);

    // Use the formula  -
    // LCE (L, R) = RMQ lcp (invSuff[R], invSuff[L]-1)
    return (RMQ(st, n, invSuff[R], invSuff[L] - 1));
}

// A function to answer queries of the longest common extension
void LCEQueries(string str, int n, Query q[],
                int m)
{
    // Build a suffix array
    vector<int> suffixArr = buildSuffixArray(str, str.length());

    // An auxiliary array to store the inverse of suffix array
    // elements. For example, if suffixArr[0] is 5, the
    // invSuff[5] would store 0.  This is used to get the next
    // suffix string from the suffix array.
    vector<int> invSuff(n, 0);

    // Build an LCP vector
    vector<int> lcp = kasai(str, suffixArr, invSuff);

    int lcpArr[n];
    // Convert to the LCP array
    for (int i = 0; i < n; i++)
        lcpArr[i] = lcp[i];

    // Build a segment tree from the LCP array
    int *st = constructST(lcpArr, n);

    for (int i = 0; i < m; i++)
    {
        int L = q[i].L;
        int R = q[i].R;

        cout << "LCE (" << L << ", " << R << ") = " << LCE(st, lcp, invSuff, n, L, R) << endl;
    }

    return;
}

// Driver Program to test above functions
int main()
{
    string str = "abbababba";
    int n = str.length();

    // LCA Queries to answer
    Query q[] = {{1, 2}, {1, 6}, {0, 5}};
    int m = sizeof(q) / sizeof(q[0]);

    LCEQueries(str, n, q, m);

    return (0);
}