Return the 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]
]


This is my approach to this question:

public class Solution {
    public ArrayList&lt;ArrayList&lt;Integer&gt;&gt; diagonal(ArrayList&lt;ArrayList&lt;Integer&gt;&gt; A) {

        //The total number of internal rows in the final list will be 2*noofrows(A) - 1
        int size = 2*A.get(0).size() - 1;
        ArrayList&lt;ArrayList&lt;Integer&gt;&gt; ans = new ArrayList&lt;ArrayList&lt;Integer&gt;&gt;();
        int count = 0;

        //Count of the number of rows
        int r = A.size();

        //Time complexity: O(size*r^2)
        //Space complexity: O(size)
        while( count &lt; size)
            {
                ArrayList&lt;Integer&gt; temp = new ArrayList&lt;&gt;();
                for( int i = 0; i &lt; r; i++ )
                {
                    for( int j = 0; j &lt; r; j++ )
                        {
                            if( ((i + j) == count))
                                temp.add(A.get(i).get(j));

                            //Continue when the indices sum up to be &gt; count
                            if((i + j) &gt; count )
                                continue;
                        }
                }
                ans.add(temp);
                count++;
            }
        return ans;
    }


I have the following questions regarding my code:

1) Is there any better approach to solve this problem?

2) How can I improve my time and space complexity?

3) Are there any redundant steps that can be removed?

Reference
I figured out a solution that improves on time complexity:
So here's the thing: the solution is actually a rearrangement of the matrix, meaning O(r^2) should be possible: visit each cell of the matrix once, determinig its location in the solution.

and it's straightforward: find the heads of the anti-diagonals, and for each one, find all its members.
finding the heads: iterate over the first row, then (when you reached end of row) iterate over the last column.
for each cell in above mentioned iteration, start gathering anti-diagonal members by advancing one row down, one column left, until either one goes beyond matrix' bounds.

few notes on the posted code:


define the variables with interface types: List&lt;List&lt;Integer&gt;&gt; ans = new ArrayList&lt;&gt;(); this includes method arguments and return value.  This will allow you to change concrete implementation with minimal changes.
use short form if (without braces) judiciously. it makes the code less clear and is prone to errors. Some ppl will say do not use this at all.  when I use it, I put the condition and action in the same line:
if( ((i + j) == count)) temp.add(A.get(i).get(j));
I think in these types of questions, where input is known and does not change, arrays are bettter suited then Collections, since they offer more concise syntax: A[i][j] is more clear then A.get(i).get(j)
naming conventions: variables start with lower case.