977. 有序数组的平方（简单）- LeetCode

题目描述

解法

双指针解法。设置一个变量index值为-1，先从左向右遍历数组，找到第一个大于等于零的元素，记录它的下标index。这一过程有三种可能结果：

• index = -1，表示没找到，说明数组全是负数
• index = 0，表示第一个数就是非负数，说明数组全是非负数
• 0 < index < 数组长度，说明数组里既有负数，也有非负数

第一、二种情况比较好办，直接按位赋值即可。第三种情况需要分别设置遍历负数的指针neg，遍历非负数的指针pos。neg向数组前端遍历，pos向数组末尾遍历，每次比较两个位置的值，把绝对值小的元素填充进结果数组，并移动指针，直至有其中一个指针到达尾部，跳出循环。跳出循环后，假设neg已经到达边界，这时只需要继续遍历pos指针即可。
时间复杂度$O(n)$，空间复杂度$O(n)$。

C++代码：

class Solution {
public:
    vector<int> sortedSquares(vector<int>& A) {
        int length = A.size();
        vector<int> ans;
        int index = -1;
        for(int i=0;i<length;i++){
            ans.push_back(-1);
        }
        for(int i=0;i<length;i++){
            if(A[i] >= 0){
                index = i;
                break;
            }
        }
        if(index == -1){
            int i = length - 1;
            for(int j=0;j<length;j++){
                ans[j] = pow(A[i],2);
                i--;
            }
        }else if(index == 0){
            for(int j=0;j<length;j++){
                ans[j] = pow(A[j],2);
            }
        }else{
            int pos = index;
            int neg = index - 1;
            int i = 0;
            for(int i=0;i<length;i++){
                A[i] = pow(A[i],2);
            }
            while(pos<length && neg>=0){
                if(A[pos] >= A[neg]){
                    ans[i] = A[neg];
                    i++;
                    neg--;
                }else{
                    ans[i] = A[pos];
                    i++;
                    pos++;
                }
            }
            if(neg >= 0){
                    for(int j=i;j<length;j++){
                        ans[j] = A[neg];
                        neg--;
                    }
                }
                if(pos<length){
                    for(int j=i;j<length;j++){
                        ans[j] = A[pos];
                        pos++;
                    }
                }
        }
        return ans;
    }
};

Python代码：

class Solution:
    def sortedSquares(self, A: List[int]) -> List[int]:
        length = len(A)
        index = -1
        ans = [-1] * length
        for i,val in enumerate(A):
            if val >= 0:
                index = i
                break
        if index == -1: # 全负
            j = length - 1
            for i,val in enumerate(A):
                ans[j] = val ** 2
                j -= 1
        elif index == 0: # 全正
            for i,val in enumerate(A):
                ans[i] = val ** 2
        else:
            pos,neg = index,index-1
            i = 0
            for j in range(length):
                A[j] = A[j] ** 2
            while pos < length and neg >= 0:
                if A[pos] >= A[neg]:
                    ans[i] = A[neg]
                    i += 1
                    neg -= 1
                else:
                    ans[i] = A[pos]
                    i += 1
                    pos += 1
            if neg >= 0:      
                for k in range(i,length):
                    ans[k] = A[neg]
                    neg -= 1 
            if pos < length:
                for k in range(i,length):
                    ans[k] = A[pos]
                    pos += 1
        return ans

官方解答

也是用双指针，不过比我的简洁多了…

class Solution(object):
    def sortedSquares(self, A):
        N = len(A)
        # i, j: negative, positive parts
        j = 0
        while j < N and A[j] < 0:
            j += 1
        i = j - 1

        ans = []
        while 0 <= i and j < N:
            if A[i]**2 < A[j]**2:
                ans.append(A[i]**2)
                i -= 1
            else:
                ans.append(A[j]**2)
                j += 1

        while i >= 0:
            ans.append(A[i]**2)
            i -= 1
        while j < N:
            ans.append(A[j]**2)
            j += 1

        return ans