Sort array with squares!

Problem Description

Given a sorted array A containing N integers both positive and negative.

You need to create another array containing the squares of all the elements in A and return it in non-decreasing order.

• Try to this in O(N) time.

Problem Constraints

1 <= N <= 10^5.
-10^3 <= A[i] <= 10^3

Input Format

First and only argument is an integer array A.

Output Format

Return a integer array as described in the problem above.

Example Input

Input 1:
 A = [-6, -3, -1, 2, 4, 5]

Input 2:
 A = [-5, -4, -2, 0, 1]

Example Output

Output 1:
 [1, 4, 9, 16, 25, 36]

Output 2:
 [0, 1, 4, 16, 25]

Solution

import Foundation

class Solution {
	func solve(_ A: inout [Int]) -> [Int] {
        var B = [Int]()
        
        var pos = A.count
        for i in 0..<A.count {
            if A[i] >= 0 {
                pos = i
                break
            }
        }
        var l = pos - 1
        var r = pos 
        
        while l >= 0 && r < A.count{
            var sl = A[l] * A[l]
            var sr = A[r] * A[r]
            if sl < sr {
               B.append(sl)
               l -= 1 
            } else if sl > sr {
                B.append(sr)
                r += 1
            } else {
                B.append(sl)
                B.append(sr)
                l -= 1
                r += 1
            }
        }
        
        if l >= 0 {
            for i in 0...l {
                B.append(A[l - i] * A[l - i])
            }
        }
        
        if r < A.count {
            for i in r..<A.count{
                B.append(A[i] * A[i])
            }
        }
        return B
	}
}

References

• InterviewBit > Sort array with squares!