977. 有序数组的平方

题目描述

输入整数数组 nums，它已经按非递减顺序排好，可能含负数。

把每个元素平方，得到一个新数组。返回这个新数组，并且它也要按非递减顺序排列。

示例 1

输入：nums = [-4,-1,0,3,10]
输出：[0,1,9,16,100]

解释：平方后是 [16,1,0,9,100]，排序得到 [0,1,9,16,100]。

示例 2

输入：nums = [-7,-3,2,3,11]
输出：[4,9,9,49,121]

数据范围

1 <= nums.length <= 10^4
-10^4 <= nums[i] <= 10^4
nums 已按非递减顺序排列。

进阶：能否用 O(n) 时间完成？

解题思路：双指针从两端向内

平方的大小只取决于绝对值。在有序数组里，绝对值最大的数一定在两端之一：最左是最小的负数，最右是最大的正数。所以用 i、j 指向两端，每次比较 |nums[i]| 和 |nums[j]|，把较大者的平方从后往前写入 ans[pos]，并把该侧指针向中间移动。pos 从 n-1 倒着走到 0，恰好填满结果。

为什么正确

ans[pos+1..] 是已排好序的最大的若干个平方；尚未处理的数恰好是 nums[i..j]，它们的平方都不超过 ans[pos+1]。

复杂度分析

时间 O(n)，每个数只处理一次；结果数组之外的额外空间 O(1)。

易错点

双指针从两端只能找到最大的平方，所以要从 ans 的末尾往前填；如果想从前往后填，就得先找到正负分界点向两边扩展。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

class Solution:
    def sortedSquares(self, nums: List[int]) -> List[int]:
        n = len(nums)
        ans = [0] * n
        i, j = 0, n - 1
        for pos in range(n - 1, -1, -1):
            if abs(nums[i]) > abs(nums[j]):
                ans[pos] = nums[i] * nums[i]
                i += 1
            else:
                ans[pos] = nums[j] * nums[j]
                j -= 1
        return ans

class Solution {
    public int[] sortedSquares(int[] nums) {
        int n = nums.length;
        int[] ans = new int[n];
        int i = 0, j = n - 1;
        for (int pos = n - 1; pos >= 0; pos--) {
            if (Math.abs(nums[i]) > Math.abs(nums[j])) {
                ans[pos] = nums[i] * nums[i];
                i++;
            } else {
                ans[pos] = nums[j] * nums[j];
                j--;
            }
        }
        return ans;
    }
}

class Solution {
public:
    vector<int> sortedSquares(vector<int>& nums) {
        int n = nums.size();
        vector<int> ans(n, 0);
        int i = 0, j = n - 1;
        for (int pos = n - 1; pos >= 0; pos--) {
            if (abs(nums[i]) > abs(nums[j])) {
                ans[pos] = nums[i] * nums[i];
                i++;
            } else {
                ans[pos] = nums[j] * nums[j];
                j--;
            }
        }
        return ans;
    }
};

func sortedSquares(nums []int) []int {
	n := len(nums)
	ans := make([]int, n)
	i, j := 0, n-1
	abs := func(x int) int {
		if x < 0 {
			return -x
		}
		return x
	}
	for pos := n - 1; pos >= 0; pos-- {
		if abs(nums[i]) > abs(nums[j]) {
			ans[pos] = nums[i] * nums[i]
			i++
		} else {
			ans[pos] = nums[j] * nums[j]
			j--
		}
	}
	return ans
}

var sortedSquares = function(nums) {
    const n = nums.length;
    const ans = new Array(n).fill(0);
    let i = 0, j = n - 1;
    for (let pos = n - 1; pos >= 0; pos--) {
        if (Math.abs(nums[i]) > Math.abs(nums[j])) {
            ans[pos] = nums[i] * nums[i];
            i++;
        } else {
            ans[pos] = nums[j] * nums[j];
            j--;
        }
    }
    return ans;
};

function sortedSquares(nums: number[]): number[] {
    const n = nums.length;
    const ans = new Array(n).fill(0);
    let i = 0, j = n - 1;
    for (let pos = n - 1; pos >= 0; pos--) {
        if (Math.abs(nums[i]) > Math.abs(nums[j])) {
            ans[pos] = nums[i] * nums[i];
            i++;
        } else {
            ans[pos] = nums[j] * nums[j];
            j--;
        }
    }
    return ans;
}

载入编辑器，自己跑一遍 →

其他解法

直接排序最简单但是 O(n log n)；分界点归并和两端双指针都是 O(n)。两端双指针从最大值开始填，不用找分界点，代码最简洁。

解法 2：平方后排序时间 O(n log n)，空间 O(log n)

把每个数平方后直接排序。

最直观，几乎不需要思考，可作为对照基准。没有利用输入已有序的性质，多了 log n 因子。

class Solution:
    def sortedSquares(self, nums: List[int]) -> List[int]:
        return sorted(x * x for x in nums)

class Solution {
    public int[] sortedSquares(int[] nums) {
        int[] ans = new int[nums.length];
        for (int i = 0; i < nums.length; i++) ans[i] = nums[i] * nums[i];
        Arrays.sort(ans);
        return ans;
    }
}

class Solution {
public:
    vector<int> sortedSquares(vector<int>& nums) {
        vector<int> ans;
        for (int x : nums) ans.push_back(x * x);
        sort(ans.begin(), ans.end());
        return ans;
    }
};

func sortedSquares(nums []int) []int {
	ans := make([]int, len(nums))
	for i, x := range nums {
		ans[i] = x * x
	}
	sort.Ints(ans)
	return ans
}

var sortedSquares = function(nums) {
    return nums.map(x => x * x).sort((a, b) => a - b);
};

function sortedSquares(nums: number[]): number[] {
    return nums.map(x => x * x).sort((a, b) => a - b);
}

解法 3：从正负分界点向两边归并时间 O(n)，空间 O(1)

先找到最后一个负数的位置 neg。负数部分平方后从右往左递增，非负部分平方后从左往右递增。用 i 从 neg 向左、j 从 neg + 1 向右，像归并两个有序数组一样每次取较小的平方放入答案。

与归并排序的合并步骤一致，便于和“合并两个有序数组”联系起来。需要先找分界点，边界情况比主解法多。

class Solution:
    def sortedSquares(self, nums: List[int]) -> List[int]:
        n = len(nums)
        neg = -1
        for i, x in enumerate(nums):
            if x < 0:
                neg = i
        ans = []
        i, j = neg, neg + 1
        while i >= 0 or j < n:
            if i < 0:
                ans.append(nums[j] * nums[j])
                j += 1
            elif j == n:
                ans.append(nums[i] * nums[i])
                i -= 1
            elif nums[i] * nums[i] < nums[j] * nums[j]:
                ans.append(nums[i] * nums[i])
                i -= 1
            else:
                ans.append(nums[j] * nums[j])
                j += 1
        return ans

class Solution {
    public int[] sortedSquares(int[] nums) {
        int n = nums.length, neg = -1;
        for (int i = 0; i < n; i++) if (nums[i] < 0) neg = i;
        int[] ans = new int[n];
        int i = neg, j = neg + 1, pos = 0;
        while (i >= 0 || j < n) {
            if (i < 0) {
                ans[pos++] = nums[j] * nums[j]; j++;
            } else if (j == n) {
                ans[pos++] = nums[i] * nums[i]; i--;
            } else if (nums[i] * nums[i] < nums[j] * nums[j]) {
                ans[pos++] = nums[i] * nums[i]; i--;
            } else {
                ans[pos++] = nums[j] * nums[j]; j++;
            }
        }
        return ans;
    }
}

class Solution {
public:
    vector<int> sortedSquares(vector<int>& nums) {
        int n = nums.size(), neg = -1;
        for (int i = 0; i < n; i++) if (nums[i] < 0) neg = i;
        vector<int> ans;
        int i = neg, j = neg + 1;
        while (i >= 0 || j < n) {
            if (i < 0) {
                ans.push_back(nums[j] * nums[j]); j++;
            } else if (j == n) {
                ans.push_back(nums[i] * nums[i]); i--;
            } else if (nums[i] * nums[i] < nums[j] * nums[j]) {
                ans.push_back(nums[i] * nums[i]); i--;
            } else {
                ans.push_back(nums[j] * nums[j]); j++;
            }
        }
        return ans;
    }
};

func sortedSquares(nums []int) []int {
	n, neg := len(nums), -1
	for i, x := range nums {
		if x < 0 {
			neg = i
		}
	}
	ans := make([]int, 0, n)
	i, j := neg, neg+1
	for i >= 0 || j < n {
		if i < 0 {
			ans = append(ans, nums[j]*nums[j])
			j++
		} else if j == n {
			ans = append(ans, nums[i]*nums[i])
			i--
		} else if nums[i]*nums[i] < nums[j]*nums[j] {
			ans = append(ans, nums[i]*nums[i])
			i--
		} else {
			ans = append(ans, nums[j]*nums[j])
			j++
		}
	}
	return ans
}

var sortedSquares = function(nums) {
    const n = nums.length;
    let neg = -1;
    for (let i = 0; i < n; i++) if (nums[i] < 0) neg = i;
    const ans = [];
    let i = neg, j = neg + 1;
    while (i >= 0 || j < n) {
        if (i < 0) {
            ans.push(nums[j] * nums[j]); j++;
        } else if (j === n) {
            ans.push(nums[i] * nums[i]); i--;
        } else if (nums[i] * nums[i] < nums[j] * nums[j]) {
            ans.push(nums[i] * nums[i]); i--;
        } else {
            ans.push(nums[j] * nums[j]); j++;
        }
    }
    return ans;
};

function sortedSquares(nums: number[]): number[] {
    const n = nums.length;
    let neg = -1;
    for (let i = 0; i < n; i++) if (nums[i] < 0) neg = i;
    const ans: number[] = [];
    let i = neg, j = neg + 1;
    while (i >= 0 || j < n) {
        if (i < 0) {
            ans.push(nums[j] * nums[j]); j++;
        } else if (j === n) {
            ans.push(nums[i] * nums[i]); i--;
        } else if (nums[i] * nums[i] < nums[j] * nums[j]) {
            ans.push(nums[i] * nums[i]); i--;
        } else {
            ans.push(nums[j] * nums[j]); j++;
        }
    }
    return ans;
}

双指针专题的其他题目

查看双指针全部题目 →