How to optimize beginner HackerEarth code

Question

I've been trying to do this HackerEarth problem but the last few test cases always seem to time out

Problem Statement: https://www.hackerearth.com/practice/algorithms/searching/linear-search/practice-problems/algorithm/joker-and-thieves-53e59f4a/

I'm a first year student, so I don't really know how to optimize that well

I tried looking at the java solution, but it just put the bigger cases into the code, effectively removing the need to optimize it

import java.io.*;
import java.util.*;
public class police {
    static int t,n,k;
    static char[][] test;
    static int max = 0;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        t = Integer.parseInt(st.nextToken());


        for (int i = 0; i < t; i++) {
           st = new StringTokenizer(br.readLine());
            n = Integer.parseInt(st.nextToken());
            k = Integer.parseInt(st.nextToken());
            test = new char[n][n];


            int ret = 0;

            for (int b = 0; b < n; b++) {
                st = new StringTokenizer(br.readLine());
                for (int a = 0; a < n; a++) {
                    test[b][a] = st.nextToken().charAt(0);
                }
            }

            for (int b = 0; b < n; b++) {
                ret += solve(test[b]); //calculate each row
            }

            System.out.println(ret);


        }    


    }
    static int solve(char[] a) { //given a row, calculate the maximum number of catches
        int ret = 0;
        for (int i = 0; i < n; i++) {
            if (a[i] == 'P') {

                for (int b = i - k; b <= i + k; b++) { //scan area to see if police can catch a thief
                    if (b >= 0 && b < n && a[b] == 'T') {
                        a[b] = 'C'; //caught
                        break;
                    }
                }

                a[i] = 'U'; //used
            }

        }
        for (int i = 0; i < n; i++) //count
            if (a[i] == 'C')
                ret++;

        return ret;
    }
}

I'm pretty sure it has something to do with the solve method, if anyone could help me that would be amazing

Answer 1

You are right, the approach you are using in your solve method is the cause of the timeout.

First you need to have an idea about algorithm complexity and Big O notation

The problem constraints are:

1 <= N <= 1000

1 <= K <= N * N

these indicate that your solution's complexity should be at most O(N * N). In other words at most two nested for loops each one with complexity O(N).

And in your solution you are doing the following:

for (int b = 0; b < n; b++) {
    ret += solve(test[b]); //calculate each row
}

OK, this loop is essential as you have to iterate over all rows in the grid. Complexity: O(N)

Then in your solve method:

for (int i = 0; i < n; i++) {
    if (a[i] == 'P') {

        for (int b = i - k; b <= i + k; b++) { //scan area to see if police can catch a thief
            if (b >= 0 && b < n && a[b] == 'T') {
                a[b] = 'C'; //caught
                break;
            }
        }

        a[i] = 'U'; //used
    }

}

These nested for loops are the real cause of the problem, the outer loop's complexity is O(N) and the inner loop complexity could be also O(N) for a high value of K. So the total complexity of the three for loops could reach O(N * N * N) which will definitely timeout.

Here's my solution to this problem, I've only modified the solve method:

static int solve(char[] a) { //given a row, calculate the maximum number of catches
        int ret = 0;
        ArrayDeque<Integer> policeMenQueue = new ArrayDeque<>(); // queue for holding positions of policemen
        ArrayDeque<Integer> thievesQueue = new ArrayDeque<>(); // queue for positions of thieves
        for (int i = 0; i < n; i++) {
            if(!policeMenQueue.isEmpty()) { // check if the leftmost policeman can catch a thief at current position (i)
                Integer mostLeftPoliceMan = policeMenQueue.getFirst();
                if(i - mostLeftPoliceMan > k) { // if he cannot then we must remove him as he will no longer be able to catch any thieves
                    policeMenQueue.removeFirst();
                }
            }
            if(!thievesQueue.isEmpty()) { // check if the leftmost thief can be caught be a policeman at current position (i)
                Integer mostLeftThief = thievesQueue.getFirst();
                if(i - mostLeftThief > k) { // if not then we must remove him as he will no longer be caught by any policemen
                    thievesQueue.removeFirst();
                }
            }
            if(a[i] == 'P') {
                if(!thievesQueue.isEmpty()) { // the leftmost thief can be caught by current policeman
                    ret++;
                    thievesQueue.removeFirst(); // ok this thief is caught, remove him
                } else {
                    policeMenQueue.addLast(i); // just add the policeman to keep his position in the queue
                }
            }
            if(a[i] == 'T') {
                if(!policeMenQueue.isEmpty()) { // the current thief can be caught by the leftmost policeman
                    ret++;
                    policeMenQueue.removeFirst(); // ok this policeman has already caught a thief (used), remove him
                } else {
                    thievesQueue.addLast(i); // just add the thief to keep his position in the queue
                }
            }
        }
        return ret;
    }

My idea: loop on each row from left to right and as the problem states: each policeman can catch only one thief and we want to maximize the number of caught thieves, so it is in our favor that each thief is caught by the leftmost policeman (as we are going from left to right).

For example consider this row:

P P T T

and imagine that K = 2

It is our favor that thief at position 3 be caught by policeman at position 1 as this policeman cannot catch the thief on position 4 and of course the policeman at position 2 can catch both thieves in this case, but remember we have to maximize the number of caught thieves and each policeman can catch only one thief, so we want that each thief be caught by the leftmost policeman who can catch him.

My solution depends on queue data structure (ArrayDeque in Java), if you don't know how it works or why i'm using it have a look here: https://www.tutorialspoint.com/data_structures_algorithms/dsa_queue.htm