977. Squares of a Sorted Array

func sortedSquares(_ A: [Int]) -> [Int] {
        var res = [Int]()

        var i = 0
        while i < A.count {
            if A[i] < 0 {
                i += 1
            }else {
                break
            }
        }

        var j = i - 1
        while j >= 0 && i < A.count{
            if A[i] > -A[j]{
                res.append(A[j] * A[j])
                j -= 1
            }else{
                res.append(A[i] * A[i])
                i += 1
            }
        }
        while j >= 0 {
            res.append(A[j] * A[j])
            j -= 1

        }

        while i < A.count {
            res.append(A[i] * A[i])
            i += 1
        }
        return res
    }