[LeetCode] 977. Squares of a Sorted Array

I. Description

문제 링크: https://leetcode.com/problems/squares-of-a-sorted-array/description/

---

II. Code

i)

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

바로 생각나는 방법으로 작성했다. Follow up 부분을 보면 O(n) 시간복잡도를 요청했기에, sorted 함수의 시간복잡도를 찾아보니 O(nlogn) 시간복잡도를 가진다는 사실을 발견했다. 따라서 다른 풀이를 생각해보았다.

 

ii)

Element를 다 넣고 정렬하는 것보다 하나씩 추가하며 적절한 위치에 넣어주는 방법은 어떨지 생각해보았다.

class Solution:
    def sortedSquares(self, nums: List[int]) -> List[int]:
        sorted_list = []
            
        for num in nums:
            square = num ** 2
            # square가 들어가야 할 인덱스 결정
            index = 0
            for i, x in enumerate(sorted_list):
                if square > x:
                    index += 1
                    
            sorted_list.insert(index, square)
            
        return sorted_list

실행해보니 Time Limit Exceeded 오류가 발생했다. 시간복잡도를 생각해보면, 첫 번째 loop에서 nums를 순회하는 것으로 O(n)이 걸리고, 그 안의 loop와 insert 작업의 시간복잡도를 O(k)로 최악의 경우, k=n이다. 따라서 총 시간복잡도는 O(1+2+3+...+n) = O(n^2)이다. 또다른 방법을 생각해보았다.

 

iii)

입력 List 자체가 정렬된 리스트라는 사실을 방금 발견했다. 그렇다면 왼쪽과 오른쪽에 포인터를 2개 놓고 좁혀가며 제곱값을 큰 순서대로 리스트 뒤쪽으로 넣어주면 어떨까 생각해보았다.

class Solution:
    def sortedSquares(self, nums: List[int]) -> List[int]:
        # 두 포인터 사용 (List는 이미 정렬된 리스트임을 활용)
        left, right = 0, len(nums) - 1
        result = [0] * len(nums)
        pos = len(nums) - 1

        while left <= right:
            left_square = nums[left] ** 2
            right_square = nums[right] ** 2
            
            if left_square > right_square:
                result[pos] = left_square
                left += 1
            else:
                result[pos] = right_square
                right -= 1
            
            pos -= 1
        
        return result

이렇게 하면 하나의 loop로 모든 과정을 처리하므로 시간복잡도가 O(n)이 나올 것이다.