반응형
문제는 여기서 볼 수 있다.
간단히 설명하자면 여러 숫자가 포함되어 있는 배열이 주어지고,
해당 배열에서 세 개의 원소의 합이 0이 되는 경우의 수를 모두 찾는 것이다.
이것도 역시 예전에 풀어봤던 문제인데, 당시에는 브루트 포스로 풀어서 느린 성능을 보였다.
그리고 역시나 Discuss에서 굉장히 좋은 성능을 보이는 코드를 찾았었다.
Discuss 코드를 보기 전에 스스로 풀어봤는데, 웬 걸!
그 코드와 같은 알고리즘을 생각해서 풀어냈다.
그런데 너무 코드가 길고 더러워서 블로그에는 예전에 찾았던 Discuss 코드를 남긴다.
(분명 Discuss 어딘가에서 봤는데, 해당 글을 도저히 찾지 못하겠다 ㅠㅠ)
바로 코드 먼저 보자.
class Solution:
def threeSum(self, nums: List[int]) -> List[List[int]]:
answer = set()
positive_nums = sorted([n for n in nums if n > 0])
positive_nums_set = set(positive_nums)
negative_nums = sorted([n for n in nums if n < 0])
negative_nums_set = set(negative_nums)
zero_nums = [n for n in nums if n == 0]
# All zero
if len(zero_nums) > 2:
answer.add((0, 0, 0))
# neg + zero + pos
if len(zero_nums) > 0:
for positive_num in positive_nums:
if -positive_num in negative_nums_set:
answer.add((-positive_num, 0, positive_num))
# neg + neg + pos
for i in range(len(negative_nums)):
for j in range(i+1, len(negative_nums)):
add = -(negative_nums[i] + negative_nums[j])
if add in positive_nums_set:
answer.add((negative_nums[i], negative_nums[j], add))
# neg + pos + pos
for i in range(len(positive_nums)):
for j in range(i+1, len(positive_nums)):
add = -(positive_nums[i] + positive_nums[j])
if add in negative_nums_set:
answer.add((positive_nums[i], positive_nums[j], add))
return list(answer)
이 알고리즘의 핵심 아이디어는 경우의 수가 딱 네 가지밖에 없다는 것이다.
위 코드의 주석에도 나와 있듯이,
1. 세 개의 수가 모두 0인 경우
2. 하나의 수가 음수, 다른 하나의 수가 0, 다른 하나의 수가 양수
3. 두 개의 수가 음수, 다른 하나의 수가 양수
4. 두 개의 수가 양수, 다른 하나의 수가 음수
딱 이 네 개 뿐이다.
그래서 각각의 경우를 나눠 구현하면, 아래와 같은 빠른 속도로 실행되는 것을 볼 수 있다.
끝!
반응형
