Squares of a Sorted Array

Problem Statement

Given an array of integers A sorted in non-decreasing order, return an array of the squares of each number, also in sorted non-decreasing order.

Questions to ask

Can i have negative numbers in my input array?

Is the input array mutable or read only?

Solution

The naive solution is to traverse the array calculating the square of each value and adding the result in a array to return. After this loop we sort the result array and return.

The time complexity of this code is O(N logn) that is the cost to sort the array.

A more optimized solution can be achived if we use two pointers and the information that the array is already sorted.

Keep one pointer in the initial position and other in the final position of the input array. Using absolute difference we discover the bigger value and we add the square of this value in the result array. The last step is to move the pointer were we found the bigger value, if was the initial we increment this pointer, otherwise we decrement the final position pointer. We are going to have a result array sorted in reverse order, reverse the array and we found a solution.

vector<int> sortedSquares(vector<int>& A) {
    vector<int> result;
    int start = 0, end = A.size()-1;
    while(start<=end){
        if(abs(A[start])<abs(A[end])){
            result.push_back(A[end]*A[end]);
            end--;
        }else{
            result.push_back(A[start]*A[start]);
            start++;
        }
    }

    reverse(result.begin(),result.end());
    return result;
}

The time complexity of this code is O(n)

question from LeetCode