Codility Training: Java Solution to TapeEquilibrium

Question: http://codility.com/demo/take-sample-test/tape_equilibrium

A non-empty zero-indexed array A consisting of N integers is given. Array A represents numbers on a tape.

Any integer P, such that 0 < P < N, splits this tape into two non-empty parts: A[0], A[1], ..., A[P − 1] and A[P], A[P + 1], ..., A[N − 1].

The difference between the two parts is the value of: |(A[0] + A[1] + ... + A[P − 1]) − (A[P] + A[P + 1] + ... + A[N − 1])|

In other words, it is the absolute difference between the sum of the first part and the sum of the second part.

Write a function:

class Solution { public int solution(int[] A); }

that, given a non-empty zero-indexed array A of N integers, returns the minimal difference that can be achieved.

Complexity:

• expected worst-case time complexity is O(N);
• expected worst-case space complexity is O(N), beyond input storage (not counting the storage required for input arguments).

The Solution

The idea is simple: first scan the array to get its sum, then scan again to compute the difference between first part and second part.

We can use add accumulatively to get the sum from a[0] to a[i], to get the difference: diff = Math.abs(2 * prefixSum - sum);

/**
 * O(1) space
 */
public int solution(int[] A) {
  if (A == null)
    throw new IllegalArgumentException("Can't be null");
  if (A.length == 0) {
    return 0;
  }
  if (A.length == 1) {
    return Math.abs(A[0]);
  }

  int sum = 0;
  for (int i = 0; i < A.length; i++) {
    sum += A[i];
  }
  int minRst = Integer.MAX_VALUE;

  // ignore last one, as second array can't be empty
  int prefixSum = 0;
  for (int i = 0; i < A.length - 1; i++) {
    prefixSum += A[i];
    int diff = Math.abs(2 * prefixSum - sum);
    if (diff < minRst) {
      minRst = diff;
    }
  }

  return minRst;
}

This is my first solution, it takes O(n) extra space.

/**
 * O(n) space
 */
public int solution1(int[] A) {
  if (A == null)
    throw new IllegalArgumentException("Can't be null");
  if (A.length == 0) {
    return 0;
  }
  if (A.length == 1) {
    return Math.abs(A[0]);
  }

  int[] prefixSum = new int[A.length];

  prefixSum[0] = A[0];
  for (int i = 1; i < A.length; i++) {
    prefixSum[i] += prefixSum[i - 1] + A[i];
  }
  int sum = prefixSum[A.length - 1];
  int minRst = Integer.MAX_VALUE;

  // ignore last one, as second array can't be empty
  for (int i = 0; i < A.length - 1; i++) {
    int diff = Math.abs(2 * prefixSum[i] - sum);
    if (diff < minRst) {
      minRst = diff;
    }
  }

  return minRst;
}

Test Code

To take a real interview test in codility, it's important to add more test cases and run them in our own IDE before submit the solution.

public static void main(String[] args) {
  TapeEquilibrium instance = new TapeEquilibrium();

  Map<int[], Integer> inAns = new HashMap<int[], Integer>();
  // the input and answer from the origin question
  inAns.put(new int[] { 3, 1, 2, 4, 3 }, 1);

  inAns.put(new int[] { 1 }, 1);
  inAns.put(new int[] { 1, 3 }, 2);
  inAns.put(new int[] { 3, 1 }, 2);
  inAns.put(new int[] { -1 }, 1);
  inAns.put(new int[] { 1, -1 }, -1);

  for (Entry<int[], Integer> entry : inAns.entrySet()) {
    assertEquals(instance.solution(entry.getKey()), entry.getValue()
        .intValue());
  }
 }

Lesson Learned

Notice common bug pattern: < or <=, init value is 0 or 1, A.length - 1 or A.length etc.
Handle and test edge cases.

Analyze time and space complexity, and whether we can improve it.