Time Limit Exceeded for CodeChef Easy problem - STRQ (Chef and Strings)

Question

I was working through a codechef easy problem here.

I've got to a solution, but it only passes the Subtask #1. For the other two subtasks, it shows TLE (Time Limit Exceeded - which is 1 second).

package com.codechef.solutions;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.PrintStream;
import java.util.StringTokenizer;

public class ChefAndStrings {

  private static final BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
  private static StringTokenizer tokenizer;

  private static String specialString = "";
  private static int numberOfQueries = 0;

  public static void main(String[] args) {
    PrintStream out = System.out;
    try {
      specialString = br.readLine();
      numberOfQueries = Integer.parseInt(br.readLine());

      for (int i = 0; i < numberOfQueries; i++) {
        String nextQuery = nextQuery();
        out.println(getGoodStringCounts(nextQuery));
      }
    } catch (IOException ex) {
      out.println("Exception occurred while reading input");
    }
  }

  private static String nextQuery() throws IOException {
    return br.readLine();
  }

  private static int getGoodStringCounts(String query) {
    tokenizer = new StringTokenizer(query, " ");
    char startLetter = tokenizer.nextToken().charAt(0);
    char endLetter = tokenizer.nextToken().charAt(0);
    int startIndex = Integer.parseInt(tokenizer.nextToken()) - 1;
    int endIndex = Integer.parseInt(tokenizer.nextToken()) - 1;

    int[] startLetterIndices = allIndicesGreaterThanMin(startLetter, startIndex);
    int[] endLetterIndices = allIndicesLessThanMax(endLetter, endIndex);

    int total = 0;
    int beginEndIndexFrom = 0;

    for (int i = 0; i < startLetterIndices.length; i++) {
      int index = startLetterIndices[i];
      for (int j = beginEndIndexFrom; j < endLetterIndices.length; j++) {
        if (endLetterIndices[j] > index) {
          beginEndIndexFrom = j;
          total += (endLetterIndices.length - beginEndIndexFrom);
          break;
        }
      }
    }

    return total;
  }

  private static int[] allIndicesGreaterThanMin(char letter, int startIndex) {
    int[] indices = new int[specialString.length() / 2];
    int currentIndex = 0;
    int index = specialString.indexOf(letter, startIndex);
    while (index != -1) {
      indices[currentIndex++] = index;
      startIndex = index + 1;
      index = specialString.indexOf(letter, startIndex);
    }
    int[] result = new int[currentIndex];
    System.arraycopy(indices, 0, result, 0, result.length);
    return result;
  }

  private static int[] allIndicesLessThanMax(char letter, int endIndex) {
    int[] indices = new int[specialString.length() / 2];
    int currentIndex = 0;
    int index = specialString.indexOf(letter);
    while (index != -1 && index <= endIndex) {
      indices[currentIndex++] = index;
      index = specialString.indexOf(letter, index + 1);
    }
    int[] result = new int[currentIndex];
    System.arraycopy(indices, 0, result, 0, result.length);
    return result;
  }
}

Initially I used List<Integer> to store the indices, but then moved to int[], after I got the error first time. Still it isn't accepting. Can you see some possible optimization there?

Answer 1

Here is an efficient algorithm that requires O(n) time for initialization and O(1) time per query.

1. Let's fix a pair (start letter, end letter)(there are only 12 such combinations). Now we can precompute the following value: count(len) - the number of strings that start with the start letter, end with the end letter and are located inside the prefix of length len. Here is a pseudo that does it:

   for start_letter <- letters:
       for end_letter <- letters:
           if start_letter != end_letter:
               start_letter_count = 0
               strings_count = 0
               for i <- 0 ... s.lenght() - 1:
                   if s[i] == start_letter:
                       start_letters_count++
                   else if s[i] == end_letter:
                       strings_count += start_letters_count
                   count(i, start_letter, end_letter) = strings_count
   

2. The answer to a (left, right, start_letter, end_letter) is count(right, start_letter, end_letter) - count(left - 1, start_letter, end_letter) - count_letter(left, right, end_letter) * count_letter(0, left - 1, start_letter), where count_letter(left, right, c) is the number of occurrences of a char c between left and right inclusively(it can be done in O(1) time with O(n) preprocessing using prefix sums). Why does this formula work? We can take all string that end before right. Now we need to subtract the number of strings that end before left. We also should subtract the number of such strings that start before left but end between left and right.

The time complexity is O(n + q) because we traverse the given string once for each pair of start and end letters, and then we use one formula that is computed in O(1) to answer each query.