[Python / BOJ 3151] 합이 0

https://www.acmicpc.net/problem/3151

3151번: 합이 0

 

문제

1. N명의 학생들의 코딩 실력 Ai가 -10000부터 10000사이의 정수로 주어진다.
2. 세 팀원의 코딩 실력의 합이 0이 되는 팀을 만들고자 한다.
3. 팀을 얼마나 많이 만들 수 있는가?

 

 

아이디어

투 포인터 활용

1. 투 포인터 활용을 위해 정렬한다.
2. 임의의 팀원 한 명을 정한다. 해당 팀원의 코딩 실력을 k 라고 하자.
3. 나머지 두 사람의 코딩 실력 합은 -k가 되어야 하며, 이 부분에서 투 포인터를 사용한다.
   • 투 포인터 합이 -k보다 크면 뒤 포인터를 감소시킨다.
   • 투 포인터 합이 -k보다 작으면 앞 포인터를 증가시킨다.
   • 투 포인터 합이 -k와 같으면 다음을 따른다.
     • p1, p2가 가리키는 값이 같으면 result에 p2 - p1 을 더한다. [1(=p1), 3, 3, 3(=p2)] 이면 +2
     • p1, p2가 가리키는 값이 다르면 result에 p2와 같은 수의 개수 만큼 더한다. [-4(=p1), 0, 2, 2, 2(=p2)] 이면 +3
     • 그 후, p1을 증가시킨다.

 

 

구현

def add_two_people(target, end):
    p1, p2 = 0, end-1
    result = 0
    tmp = N
    while p1 < p2:
        hap = IN[p1] + IN[p2]
        # hap이 더 크면 줄여야 한다.
        if hap > target:
            p2 -= 1
        # hap이 더 작으면 키워야 한다.
        elif hap < target:
            p1 += 1
        # target 이면 카운트
        else:
            # p1, p2가 가리키는 값이 같을 때
            if IN[p1] == IN[p2]:
                result += p2 - p1
                
            # p1, p2가 가리키는 값이 다를 때
            else:
                # 이 if 문이 있어야 시간이 줄어든다
                if tmp > p2:
                    tmp = p2
                    while tmp >= 0 and IN[p2] == IN[tmp-1]:
                        tmp -= 1
                result += p2 - tmp + 1
            p1 += 1

    return result


N = int(input())
IN = sorted(list(map(int, input().split())))

# 3명 미만이면 불가
if N < 3:
    print(0)
    exit()

result = 0
for p3 in range(N-1, 1, -1):
    result += add_two_people(-IN[p3], p3)

print(result)

 

 

다른 사람 코드

• 스위핑, 조합론 출처

n = int(input())
arr = list(map(int, input().split()))

dp = [0] * 40001
total = 0

# 배열의 인덱스가 음수가 되지 않도록 적당한 숫자 20000를 더해준다.
for i in range(n):
    total += dp[20000 - arr[i]]

    for j in range(i):
        dp[20000 + arr[j] + arr[i]] += 1

print(total)