[BOJ, Python] 백준 3151번_합이 0

박상민·2024년 10월 5일
0

Algorithm

목록 보기
15/20
post-thumbnail

백준 3151번

우선 문제는 투포인터 문제이다.
하지만 학생 3명 점수의 총합을 찾아야 된다는 조건이 있다. 삼포인터 같은 알고리즘은 존재하지 않는데 문제를 어떻게 풀어야 할까?

나는 아래처럼 문제를 풀고자 했다.

풀이 이론

  1. 점수 배열을 오름차순으로 정렬한다.
  2. 학생 한 명을 고정한다.
  3. 나머지 학생 2명의 점수 총합을 투포인터 알고리즘으로 탐색한다.
  • 이처럼 학생 한 명을 고정한다면 투포인터로 문제를 풀 수 있다.

학생 한 명을 고정한다는 것을 제외하면 일반적인 투포인터 문제와 비슷하기에 코드는 금방 구현했다. 그러나 두번째 문제가 발생했다.
문제 발생 코드

import sys
input = lambda: sys.stdin.readline().rstrip()

N = int(input())
A = list(map(int, input().split()))

# 0~(양수가 되는 지점 전)까지의 각 점을 팀원 중 한 명으로 삼고 투 포인터 방식으로 전개
A.sort()

S = N-2
for i in range(N):
    if A[i] > 0:
        S = i
        break
        
ans = 0
for i in range(0,S):
    tmp = A[i]
    start,end = i+1, N-1
    while start < end:
        s = tmp + A[start] + A[end]

        if s==0:
            if A[start] == A[end]:
                ans += end - start
            ans += 1
            start += 1 ## 문제 발생 부분

        elif s < 0: # s가 0보다 작다면 값을 키워야함
            start += 1
        else: # s가 0보다 크다면 값을 작게 해야됨
            end -= 1

print(ans)

주어진 테스트 코드는 모두 통과하는데 코드를 제출하면 1%에서 실패했다.
처음에는 한참동안 무엇이 문제인지 찾지 못했다.

문제점
현재 총합(s)의 값이 0이라면 시작 지점(start)를 +1하고 있다. 바로 이 부분이 문제이다.

이해를 돕기 위해 반례 하나를 보자

반례
입력
6
-8 3 3 5 5 5
출력
6

위 반례의 경우 나의 코드가 진행되는 흐름을 살펴보자.

  1. 점수가 -8인 학생을 고정한다.
  2. start의 인덱스가 1, end의 인덱스가 N-1(5)인 경우 총합은 0이다.
  3. ans에 +1을 하고 start의 인덱스를 2로 증가시킨다.
  4. start의 인덱스가 2, end의 인덱스가 N-1(5)인 경우 총합은 0이다.
  5. ans에 +1을 하고 start의 인덱스를 3로 증가시킨다.
  6. 이후에는 총합이 0이 되는 경우가 없다.

이제 문제점이 명확히 보인다.
start의 인덱스가 1인 경우 end의 인덱스 3,4,5의 경우 모두 총합이 0이 된다.
그러나 처음 총합이 0이 되는 순간 start의 인덱스를 +1하기 때문에 end 인덱스 3,4의 경우는 찾지 못하고 스킵이 되는 것이다. start 인덱스가 2인 경우도 마찬가지다.

해당 문제점은 첫 제출 후 바로 발견했다.

그렇다면 총합이 0이 되는 경우 start, end의 인덱스를 어떻게 변경해야 할까?
start를 옮기는 것도, end를 옮기는 것도 모두 문제가 된다.
이 문제를 해결하고자 1시간 동안 고민을 하고 다양한 시도를 했지만 결국 해결하지 못했다.

결국 다른 분의 풀이를 참고해서 문제를 해결했다.
풀이를 보고 문제를 보는 시선이 너무 좁았음을 깨닫고 반성했다.

해결법

  • 배열 중 같은 숫자가 몇 개 있는지 파악한 후 총합이 0이 되는 경우 배열 중 end 위치의 숫자의 갯수를 ans에 더해준다.
  • 이렇게 된다면 탐색하지 탐색하지 못하는 경우가 생기지 않는다.
  • 배열은 정렬되어 있기 때문에 end 인덱스의 값보다 작은 값들의 총합이 항상 0보다 작기에 탐색할 필요가 없다.

나는 같은 숫자의 갯수를 딕셔너리로 저장해서 사용했다.
그러나 collections의 Counter 함수를 사용하면 더 간단하게 풀 수 있다.

from collections import Counter
A = 배열
cnt_ = Counter(A)

풀이 코드

import sys
input = lambda: sys.stdin.readline().rstrip()

N = int(input())
A = list(map(int, input().split()))

# 0~(양수가 되는 지점 전)까지의 각 점을 팀원 중 한 명으로 삼고 투 포인터 방식으로 전개
A.sort()

S = N-2
for i in range(N):
    if A[i] > 0:
        S = i
        break

dict = {}
for a in A:
    if a in dict:
        dict[a] += 1
    else:
        dict[a] = 1

ans = 0
for i in range(0,S):
    tmp = A[i]
    start,end = i+1, N-1
    while start < end:
        s = tmp + A[start] + A[end]

        if s==0:
            if A[start] == A[end]:
                ans += end - start
            else:
                ans += dict[A[end]]
            start += 1

        elif s < 0: # s가 0보다 작다면 값을 키워야함
            start += 1
        else: # s가 0보다 크다면 값을 작게 해야됨
            end -= 1

print(ans)

결과

0개의 댓글