Sort squares is to sort the square of elements in a sorted array. Sounds easy, right? Since it uses sorting algorithms, the time complexity is O(n^2). Now you are asked to optimize it. The complexity should be O(n).

How to optimize? We can use the “merge two sorted arrays” approach. We divide the sorted array into half. The first half has all the negative numbers. The second half has all the numbers that are equal or larger than zero. Then we merge the two arrays of the squares of the original numbers. This solution has the time complexity to O(n).

You can see that sort squares is a “merge two sorted arrays” question under disguise.

**VMwave Interview Question:**

Input: sorted integers: example: [-2, 1, 2, 3]

Output: sorted square if the input: [1, 4, 4 ,9]

**Code:**

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
import java.util.Arrays; public class SortSquares { //Two pointers, Time O(n), Space O(n), n is array length public static int[] sortSquares(int[] a) { int n = a.length; int[] res = new int[n]; int i = 0, j = n - 1; for (int k = n-1; k >= 0; k--) { if (Math.abs(a[i]) > Math.abs(a[j])) { res[k] = a[i] * a[i]; i++; } else { res[k] = a[j] * a[j]; j--; } } return res; } public static void main(String[] args) { int[] a = {-2, 1, 2, 3}; System.out.print("Input: "); System.out.println(Arrays.toString(a)); int[] b = sortSquares(a); System.out.print("Sort squares: "); System.out.println(Arrays.toString(b)); } } |

**Output:**

Input: [-2, 1, 2, 3]

Sort squares: [1, 4, 4, 9]

**O Notation:**

Time complexity: O(n)

Space complexity: O(n)

**Note:**

If you have any questions or want to put comments, please post at youtube. I will answer you!

The code displayed here is the simplified version after the enhancement in the 2nd edition of Java coding interview pocket book.

**Actionable:**

Download SortSquares.java

Sort squares with optimization tutorial

Merge two sorted arrays (2 solutions)