Solving UVa 978 With Unit Tests


Last week, I wrote about how unit testing can improve your competitive programming practice. As promised, this week I’m going to go through an example of the unit testing process. For an example problem, I’ll be using UVa 978: Lemmings Battle! See above for a picture of a ferocious lemming.

Problem Summary

Lemmings Battle! is rated Level 4 on uHunt, but I found it to be more of a Level 3 problem. It can be solved by simulating the game rules presented in the problem statement, and I didn’t run into any issues related to runtime performance. Other than getting the details of the simulation right, the only potential challenge is what primary data structure to use. uHunt suggests a TreeSet, but I found a TreeMap to be more appropriate.

Here are the key points of this problem:

  • The goal is to simulate the outcome of a series of battles, and report the final result.
  • The warring parties are the green lemming army and the blue lemming army.
  • The starting conditions consist of a fixed number of battlefields and a list of lemmings on each side. Each lemming starts with a power level (an integer value).
  • At the beginning of each round, a pair of lemmings, one from each army, is assigned to each battlefield in descending order of power level.
  • During the round, each lemming pair fights. In each fight, the lemming with the higher power wins, but its power is reduced by an amount equal to the other lemming’s power. If the lemmings in the pair have equal power, both lemmings end the round with 0 power.
  • At the end of the round, lemmings with nonzero power are returned to their armies.
  • Rounds continue as long as both armies have lemmings available to fight.
  • The problem output must identify the winning army (if any) and the power values of surviving lemmings, in descending order.

TreeMap and TreeSet

In Java, the appropriate data structure to use for this problem is TreeMap or TreeSet. These are classes that both provide ways to efficiently store and retrieve objects in sorted order. You can look up the full list of similarities and differences between the two classes, but for the purpose of this problem, the important differences are:

  • TreeMap maps a key to its corresponding value. If you want to update or retrieve a single value, you have to know its key. With TreeSet, you just store the values. There are no keys, and therefore no way to look up a value directly, though you can retrieve the values in other ways, such as iterating through all of them or retrieving the first or last value.
  • TreeSet does not allow duplicate values. If you want to add an object to a TreeSet that matches one that it already contains, you need to remove the existing one first. Similarly, TreeMap doesn’t allow duplicate keys, since the key uniquely identifies the object to store or retrieve. However, objects that are equal to each other can be associated with multiple keys in a TreeMap.

We’ll use this information later on to pick the appropriate data structure to use to solve this problem.


JUnit is a popular Java unit testing tool, and I’ll be using it for this example. To “install” JUnit, you just add its JAR file to your classpath. Links to the required JARs are available on the JUnit GitHub site.

Having JUnit in your classpath allows you to compile and run your test class, which we’ll get to shortly.

A Testable Solution Template

Having a template that you use to start each puzzle solution helps avoid repetitive typing. But a template that isn’t written with unit testing in mind isn’t likely to lead to a testable solution. To that end, I adjusted my standard template so that it works well with tests while also supporting the stdin/stdout model required by UVa Online Judge. That allows me to run tests locally and then submit the same code to UVa OJ without making any changes.

My updated template contains the testable code, and a new template contains a sample test.

The following aspects of the updated make it testable:

  • Input and output implementations are passed to a setupIO method rather than being hard-coded in the main method. This allows the Main class to use stdin for input by default, while giving the test class the option to override this implementation with a FileReader that reads from a test file.
  • Since the public static void main method is bypassed by the test class, it isn’t covered by tests. Therefore, it contains only the minimum code required to kick off the rest of the program. That code comes directly from the template, so there’s less chance of adding bugs to it in the course of solving a problem.
  • Most methods and instance variables are public so they can be called and inspected by the test class.

The other half of the testable template is, which contains the test code. Here are the key parts of that class:

  • The setUp method runs before each test method. It creates an instance of Main and sets up I/O to read the test input file.
  • Methods marked @Test are run by the test runner. They are run in arbitrary order, so they must be independent of each other.
  • Calling assertEquals is a simple way to verify that the result of a method call or the contents of an instance variable is what you expect it to be at a point in the test.

To run the tests in MainTest against the code in Main, use the following commands.

java org.junit.runner.JUnitCore MainTest

If all goes well and your classpath is set correctly, you should get a result like this:

JUnit version 4.10
Time: 0.006

OK (1 test)

If any tests fail, you’ll get a stack trace indicating where the failure occurred.

Testing and Implementing

If we’re going to use unit testing for programming puzzles, we might as well consider using test-driven development (TDD).

Here’s how you could use TDD during the implementation step of a puzzle-solving process:

  • Based on your pseudocode, select a small part of the solution to implement.
  • Write test data and test code that together will verify this part of your implementation.
  • Run all of your tests, and ensure that they all pass except for the new one.
  • Write solution code that will make the new test pass.
  • Run all of your tests and ensure that they all pass.

If you repeat these steps as you work towards a complete solution, you should find yourself spending less time debugging at the end, since you’ll be verifying each part of your solution as you go.

Solving UVa 978

As an example of how TDD can work for programming puzzles, here’s a sequence of tests and implementation that results in a solution to UVa 978.

Process parameters

The template contains pre-tested functionality that we can use for free, including code that retrieves the number of test cases. So the first code we need to test is a method that retrieves the problem’s three parameters:

  • The number of battlegrounds
  • The number of green lemmings
  • The number of blue lemmings

For a puzzle solution, there are some advantages to declaring variables for these values at the class instance level. This also allows our test to do the following:

  • Call a processParams function.
  • Verify that all of the instance variables have the correct values, based on the numbers in the test data file.

For this first test, the actual values used in the test data aren’t too important, so we can use any reasonable numbers, or just use the sample input. Future tests will depend more on carefully-selected input.

With TDD, the idea is to start with a failing test, and write the minimal amount of code required to make it pass. How granular your tests should be, and therefore how little code you write between test runs, is up to you. For example, a compile error is considered a test failure. So if your test code calls processParams but you haven’t written a processParams function, you’ll get a cannot find symbol error from the compiler. This will prompt you to write that function. Whether you actually do the compile step, or just think about doing it, is your choice.

Even if you don’t run your test before you implement processParams, you should definitely run it after you create the method and the instance variables, but before you initialize them. That way you can verify that your test informs you when these variables are not set correctly. (They’ll be set to their default value of 0). It’s important to verify that your test can fail, since otherwise you might inadvertently create a test that can never fail, and therefore isn’t useful for verifying correctness.

Once you verify that your test fails, you can implement processParams and run your test again to verify that it passes.

Build armies

After the three initial parameters, the remainder of the input for each test case consists of a list of power levels for each lemming in the green and blue armies.

To store a list of integer values, we’ll need some kind of collection. It’s clear from the sample input that there can be multiple lemmings with the same power level. And it’s clear from the problem statement that the power value is the one we need to sort on, since the rules require us to use the lemming with the highest power as we’re assigning lemmings to battlefields.

For both TreeSet and TreeMap, duplicate values require special handling of some kind. My approach to handling duplicate power values is to use a second field to store the number of lemmings with a given power. If the collection contains a lemming with a particular power, and another lemming with the same power needs to be stored, I can increment the counter field. If there are multiple lemmings with the same power value and I remove one, I decrement the counter. I could do that with a TreeSet<Pair>, where Pair is a custom object containing two integer values (power and count). But updating the count would take an extra remove step, because Set.add() does nothing if the entry being added matches (according to compareTo) an existing entry. Therefore, I decided that TreeMap<Integer, Integer> is the more appropriate data structure. The first integer is the power level and the second integer is the count. Since two integers are exactly what we need to store, this also avoids having to implement a custom class with a custom compareTo method.

To encapsulate the counting logic required to use the TreeMap for storing multiple identical power values, I decided to implement a function putPower whose responsibility is to add a lemming to an army. To test that function, we can do the following:

  • Pass in a power value and verify that the count is 1.
  • Pass in the same power value again, and verify that the count is 2.
  • Run the test and verify that it fails.

You could also write other tests to verify other cases.

Then to implement the function:

  • Retrieve an entry from the army TreeMap using the power as a key.
  • If get returns null, then put the power with a count of 1.
  • Otherwise, increment the previous count by one.

Now that we have a way to add a lemming to an army, it’s easy to test and implement a buildArmies function to initialize the armies. To test it, just call it and verify that the largest power in each army is as expected based on the test data. To implement it, read the correct number of lines (based on the test case parameters) for each army and call putPower on each one.

Assign battlefields

At the beginning of each round, we need to assign lemmings to battlefields. To do that, we need a way to get the next lemming from an army, while following the lemming ordering rules. I created a getNextLemming function for that purpose. The rule from the problem statement is that the lemming with the highest power is always the next one to fight. So our test must verify that this is in fact the lemming that is retrieved. To implement getNextLemming, we can use the fact that TreeMap keys come out of the tree in sorted order. This is of course why we chose this data structure in the first place. TreeMap.pollLastEntry() removes and returns the lemming with the highest power value. The remaining step is to keep the count field up to date in cases where there are multiple lemmings with the highest power value.

Now that we have a way to get the correct lemming from an army, we can assign a lemming to each battlefield. A test of assignBattlefields is to verify that the highest-powered lemmings is assigned to the lowest-numbered battlefield. Implementing the function is just a matter of looping through each battlefield, and assigning to it the lemming returned by getNextLemming.


Once battlefields are assigned, the lemmings can fight. To verify the result of a fight, we just need to verify that the power levels on a battlefield follow the problem rules: the power value of the higher-powered lemming drops by an amount equal to the power value of the lower-powered lemming, or both drop to 0 if the lemmings have equal power. To implement a fight round, we just loop through all of the battlefields and carry out those subtractions.

Send back

After the fights, surviving lemmings are sent back to their armies. To verify that this works correctly, we can check the power values and counts. To implement sendBackToArmies, we can re-use our putPower function: for each battlefield with a lemming of power greater than 0, call putPower for the appropriate army.


We now have almost everything written and tested, and we just need to string it together: For each test case, we call processParams and buildArmies, then repeatedly call assignBattlefields, fight, and sendBackToArmies until one or both armies are empty.


The only remaining task is to print the results. The strings to be printed are listed in the problem statement, and are entirely based on the size of each army. The winning team has a nonzero army size, and the losing team has an empty army. For the winning team, we have to print the remaining lemming powers in descending order. This is just a matter of calling getNextLemming until it returns null.

We could test the printResults and printArmy functions using string comparisons, but it’s probably better just to diff the results with the sample output and with the expected results from tests cases found on uDebug.

I put the complete solution and tests as described on GitHub.

Thoughts on Implementing and Testing UVa 978

UVa 978 is a reasonable contest problem to use when trying out test-driven development:

  • It’s not trivial to solve. Both the on-paper solution and the implementation contain issues to think about.
  • It naturally breaks down into multiple short functions, each of which can be tested independently.
  • The natural solution breakdown doesn’t require any tweaks to fix runtime performance, as long as you use an appropriate tree for the main data structure. In other words, you don’t have to sacrifice design for performance.

The unit testing approach certainly presents unique challenges compared to a traditional puzzle-solving process. It may result in more code and more time spent, though it’s also possible that the tests lead to a concise solution and the careful approach results in less time spent debugging. But ultimately time and code size aren’t the main concern for beginning and intermediate students. Using unit tests for puzzle solutions is worthwhile if it results in a better understanding of the solution, and the problem-solving process. If you eventually get to the point where you’re optimizing your performance in live contests, you’ll probably drop the unit testing approach. But the experience you get writing solutions in this way will give you a valuable perspective on how to write the best solution you can.

(Image credit: Sander van der Wel)