Anti Diagonals

Give a N*N square matrix, return an array of its anti-diagonals. Look at the example for more details.

Example:

		
Input: 	

1 2 3
4 5 6
7 8 9

Return the following :

[ 
  [1],
  [2, 4],
  [3, 5, 7],
  [6, 8],
  [9]
]

Input : 
1 2
3 4

Return the following  : 

[
  [1],
  [2, 3],
  [4]
]

public class Solution {
    public int[][] diagonal(int[][] A) {
        int n = A.length;
        int[][] ans = new int[2 * n - 1][];
        int size = 1;
        for (int i = 0; i < 2 * n - 1; i++) {
            ans[i] = new int[size];
            if (i >= (n - 1))
                size--;
            else
                size++;
        }
        int index[] = new int[2 * n - 1];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                ans[i + j][index[i + j]] = A[i][j];
                index[i + j] += 1;
            }
        }
        return ans;
    }
}