USACO Guide: Counting Haybales

This post marks the beginning of what I hope will become a long series of posts detailing my solutions to many USACO problems. USACO, or the USA Computing Olympiad, is an annual online computer programming competition. Each year, USACO holds four competitions, each consisting of 3-4 problems to be solved in 3-5 hours. There are four levels of competition-Bronze, Silver, Gold, and Platinum-and five accepted languages-C, C++, Java, Pascal, and Python. Throughout my posts, I will mainly be covering previous Bronze and Silver problems and using Java 8, the only allowed version of Java, or occasionally Python. Every language has its own restrictions for run time and memory usage, and every test case has the same value to one’s overall competition score.

The Problem

The problem text is shown below (USACO Page):

Problem Text (http://www.usaco.org/index.php?page=viewproblem2&cpid=666)

My Solution

import java.io.*;
import java.util.*;

public class CountingHaybalesDecember2016 {
    private static String fileNamePrefix = "haybales";
    private static int[] haybales;

    public static void main(String[] args) throws IOException {
        BufferedReader reader = new BufferedReader(new FileReader(fileNamePrefix + ".in"));
        StringTokenizer stringTokenizer = new StringTokenizer(reader.readLine());

        int numHaybales = Integer.parseInt(stringTokenizer.nextToken()), numQueries = Integer.parseInt(stringTokenizer.nextToken());
        haybales = new int[numHaybales + 1];

        stringTokenizer = new StringTokenizer(reader.readLine());

        for (int i = 0; i < numHaybales; i++) {
            haybales[i] = Integer.parseInt(stringTokenizer.nextToken());
        }

        haybales[numHaybales] = -1; // To have an element before all others

        Arrays.sort(haybales);

        int[][] queries = new int[numQueries][2];

        for (int i = 0; i < numQueries; i++) {
            stringTokenizer = new StringTokenizer(reader.readLine());

            queries[i][0] = Integer.parseInt(stringTokenizer.nextToken());
            queries[i][1] = Integer.parseInt(stringTokenizer.nextToken());
        }
        reader.close();

        TreeMap<Integer, Integer> haybaleIndexMap = new TreeMap<Integer, Integer>();

        for (int i = 0; i < haybales.length; i++)
            haybaleIndexMap.put(haybales[i], i);

        PrintWriter writer = new PrintWriter(new BufferedWriter(new FileWriter(fileNamePrefix + ".out")));

        for (int[] query : queries) {
            int start = query[0], end = query[1];

            // Find the index of the first item less than or equal to the start and the end and subtract them
            int startKey = haybaleIndexMap.floorKey(start - 1), endKey = haybaleIndexMap.floorKey(end);
            writer.println(haybaleIndexMap.get(endKey) - haybaleIndexMap.get(startKey));
        }

        writer.close();
    }
}

Explanation

The first lines of my code are rather simple: they just open the file with a BufferedReader and use a StringTokenizer to parse the data in the first 2 lines (the number of hay bales, the number of queries, and the positions of each hay bale). I also initialize haybales, an array of ints with a length that is 1 greater than the number of hay bales. The reasoning for this length will be discussed later. Finally, the first numHaybales elements of haybales are set to contain the position of a different hay bale, leaving the final item (haybales[numHaybales]) unset.

BufferedReader reader = new BufferedReader(new FileReader(fileNamePrefix + ".in"));
StringTokenizer stringTokenizer = new StringTokenizer(reader.readLine());

int numHaybales = Integer.parseInt(stringTokenizer.nextToken()), numQueries = Integer.parseInt(stringTokenizer.nextToken());
haybales = new int[numHaybales + 1];

stringTokenizer = new StringTokenizer(reader.readLine());

for (int i = 0; i < numHaybales; i++) {
    haybales[i] = Integer.parseInt(stringTokenizer.nextToken());
}

Next, this final element is set to -1 and haybales is sorted. This results in an array in which haybales[0] == -1 and the remaining numHaybales elements are the positions of every hay bale, in ascending order.

haybales[numHaybales] = -1; // To have an element before all others

Arrays.sort(haybales);

This is why I set the length of haybales to one greater than numHaybales, to include this -1 at the beginning. It will be very important to later code.

Next, I initialize a two-dimensional array queries from the remaining lines of the input file. Each element of queries is an array of ints with length 2, where the first element of that array is the lower bound of the query (A) and the second is the upper bound (B).

int[][] queries = new int[numQueries][2];

for (int i = 0; i < numQueries; i++) {
    stringTokenizer = new StringTokenizer(reader.readLine());

    queries[i][0] = Integer.parseInt(stringTokenizer.nextToken());
    queries[i][1] = Integer.parseInt(stringTokenizer.nextToken());
}

reader.close();

The following code initializes haybaleIndexMap, a TreeMap of Integers that is then given the position of each element of haybale as a key and its index as a value (including the -1 at index 0).

TreeMap<Integer, Integer> haybaleIndexMap = new TreeMap<Integer, Integer>();

for (int i = 0; i < haybales.length; i++)
    haybaleIndexMap.put(haybales[i], i);

Finally, the actual solution is obtained and logged to a file through a PrintWriter. For each query, the lower and upper bounds are retrieved and stored in start and end, respectively. Then, the function of TreeMaps that are essential to the solution is used, floorKey (JavaDoc link). By using this method, we are able to obtain the largest key (position) less than or equal to each of the bounds. Subtracting the values of these bounds gives the difference between the indexes of each position in the sorted haybales array, meaning that it will result in the number of hay bales between the two. That value is the solution and is printed to the file.

PrintWriter writer = new PrintWriter(new BufferedWriter(new FileWriter(fileNamePrefix + ".out")));

for (int[] query : queries) {
    int start = query[0], end = query[1];

    // Find the index of the first item less than or equal to the start and the end and subtract them
    int startKey = haybaleIndexMap.floorKey(start - 1), endKey = haybaleIndexMap.floorKey(end);
    writer.println(haybaleIndexMap.get(endKey) - haybaleIndexMap.get(startKey));
}

writer.close();

This is the reason that the -1 was included: there will always be a value less than or equal to any hay bale position, so queries involving the hay bale at the smallest position (the hay bale at haybales[1]) will not result in a null return value from floorKey and the correct answer will still be obtained.

Results

This solution successfully passes all test cases without exceeding the memory or runtime limits, as seen below:

Successful test cases