Merge k sorted lists in Java

Question

I hava a question similiar with mege in Python but I am trying to write Java code. Here is my code: The input is: The first line - a number of arrays (k); Each next line - the first number is the array size, next numbers are elements.

Max k is 1024. Max array size is 10*k. All numbers between 0 and 100. Memory limit - 10MB, time limit - 1s. Recommended complexity is k ⋅ log(k) ⋅ n, where n is an array length.

Example input:

4            
6 2 26 64 88 96 96
4 8 20 65 86
7 1 4 16 42 58 61 69
1 84

Example output:

1 2 4 8 16 20 26 42 58 61 64 65 69 84 86 88 96 96 

I get memory limit errors. I use System.gc() to clear the memory but I recieve timeout error

Here is my code:

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
public class TaskFSolution {   
    public static void main(String[] args) throws IOException {
        BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));

        int k = Integer.valueOf(reader.readLine());

        short[][] arrays = new short[k][];

        int arrayLength = 0;
        int[] minPositionInArrays = new int[k];
        for (int i = 0; i < k; i++) {
            String str = reader.readLine();
            String[] elements = str.split(" ");
            arrays[i] = new short[Short.valueOf(elements[0]) + 1];
            for (int j = 0; j < elements.length; j++) {
                arrays[i][j] = Short.valueOf(elements[j]);
            }
            arrayLength += Short.valueOf(elements[0]);
            minPositionInArrays[i] = 1;

            if (arrayLength % 100 == 0) {
                System.gc();
            }
        }

        reader.close();

        StringBuilder sb = new StringBuilder();
        int currentArrayNumber = 0;
        int currentPositionInArray = minPositionInArrays[currentArrayNumber];
        while (arrayLength > 0) {
            short currentMinValue = 101;
            for (int i = 0; i < k; i++) {
                if (minPositionInArrays[i] <= arrays[i][0] && currentMinValue >= arrays[i][minPositionInArrays[i]]) {
                    currentMinValue = arrays[i][minPositionInArrays[i]];
                    currentArrayNumber = i;
                    currentPositionInArray = minPositionInArrays[i] + 1;
                }
            }
            sb.append(currentMinValue + " ");
            minPositionInArrays[currentArrayNumber] = currentPositionInArray;
            arrayLength--;
        }

        System.out.println(sb.toString().trim());
    }
}

Answer 1

Are you getting the correct output with the recommended complexity? It's been a long time since I've written my own sort but start by confirming that before any of my recommendations below.

I suspect str.split(" ") is causing problems for you because you're creating lots of strings which you're about to parse and then throw away. How about parsing the integers directly from str in a single pass instead?

Try changing sb.append(currentMinValue + " ") to sb.append(currentMinValue).append(" "). It might be a small optimization but it's easy to do.

How about splitting your main into two methods? The first for everything that uses reader, the second for everything that uses sb. Besides improving the structure of the program this provides a natural point for garbage collection.

What command-line args are you giving to the JVM? Since you know your memory limit try allocating it all at the start to save time allocating more memory later. You might also play with different garbage collection strategies. As you're discovering, forcing garbage collection isn't a good idea. If you're correctly setting the maximum in the JVM it should automatically garbage collect as needed to prevent running out of memory.

Hope at least one of these suggestions helps. I'm curious what gives the biggest improvement.

Answer 2

I was curious about why you kept the initial size inside each of the arrays until I finaly noticed the
arrays[i][0] in the check if that array still has elements to proces during the algorithm. If we replace that with arrays[i].length-1 we no longer need to store those.

Another major issue, like Nathan pointed out is how many String objects you're creating while parsing the input. The easiest solution to this is to use a Scanner instead of a BufferedReader. That way, you can directly read the nextShort().

    Scanner reader = new Scanner(new InputStreamReader(System.in));

    short k = reader.nextShort();
    short[][] arrays = new short[k][];

    int arrayLength = 0;
    int[] minPositionInArrays = new int[k]; // note here: now correctly initialised to 0's
    for (int i = 0; i < k; i++) {
        short size = reader.nextShort();
        arrays[i] = new short[size];
        for (int j = 0; j < size; j++) {
            arrays[i][j] = reader.nextShort();
        }
        arrayLength += size;
    }
    reader.close();

---

The recommended complexity is O(k ⋅ log(k) ⋅ length_array). Yours seems to be O(k ⋅ (k ⋅ length_array)).

Since we need to proces k ⋅ length_array items anyway, we can't change that. So how can you get the other k in that complexity formula down to a log(k)?

The solution is to sort all the lists initially and then use the fact that they're sorted when processing each element. How to actually implement this I'll leave up to you :)

---

Final remark: If the memory limit is 10MB and the input is at most 1024*(10*1024) = 10.485.760 numbers (~=10MB) this seems almost impossible to me. That would mean storying all numbers into a single array in the first place and sorting them in-place. But you don't know how big this array needs to be before you read the last line of the input.

I'm assuming the actual limit is a bit more lenient, or it wasn't set with java in mind :)