[백준/투 포인터] - 합이 0인 네 정수

문제 링크
✅ PyPy3로 제출했음.

코드

n = int(input())
_list = [[] for _ in range(4)] # [A], [B], [C], [D]가 포함된 리스트

# 각 리스트에 값 넣기
for _ in range(n):
    _input = list(map(int, input().split()))
    idx = 0
    for i in range(4):
        _list[i].append(_input[idx])
        idx += 1

# A[a] + B[b] + C[c] + D[d] = 0
# A[a] + B[b] = -(C[c] + D[d]) => a와 b의 합, c와 d의 합을 리스트로 변환   
ab, cd = [], []

for i in range(n):
    for j in range(n):
        ab.append(_list[0][i] + _list[1][j])
        cd.append(_list[2][i] + _list[3][j])

# 합이 0이 되는 쌍을 찾으려면, 정렬을 해야 한다.
# Why? 합이 0보다 큰지 작은지에 따라 포인터를 움직여야 하는데, 정렬하지 않으면 의미가 없음.(=동작 안함)
ab.sort()
cd.sort()

l, r = 0, len(cd) - 1
cnt = 0 # 총 쌍의 개수

while l < len(ab) and r >= 0:
    ab_sum = ab[l]
    cd_sum = cd[r]
    total_sum = ab_sum + cd_sum # 합
    dup_ab, dup_cd = 0, 0
    # 합이 0인 쌍이라면
    if total_sum == 0:
        # 중복되는 수 찾기
        while l < len(ab) and ab_sum == ab[l]:
            dup_ab += 1
            l += 1
        while r >= 0 and cd_sum == cd[r]:
            dup_cd += 1
            r -= 1
        cnt += dup_ab * dup_cd # 총 경우의 수 갱신
    elif total_sum > 0: # r을 줄여야 한다.
        r -= 1
    else: # l을 늘려야 한다.
        l += 1

print(cnt)

풀이

문제의 목표

• A[a], B[b], C[c], D[d]의 합이 0인 (a, b, c, d) 쌍의 개수 구하기

예전에 풀었던 문제이다. 4개의 포인터를 다루는 방식 말고 다른 것을 사용했다.

로직은 다음과 같다.(전제: (a,b), (c,d)를 정렬)

• (a, b), (c, d)를 묶어서 각 원소들의 합을 리스트로 만든다.
• 각 포인터를 l=0, r=len(cd)-1로 설정하고 투 포인터를 진행한다.
• 만약, 두 합이 0이라면 중복되는 수가 있을 수 있으므로 찾아준다.((a,b), (c,d) 모두에 대해서 찾는다.)
  • 중복되는 수가 있다면, 두 개수((a,b), (c,d)에 대한 개수)를 곱해준다.
• 만약 합이 0보다 크다면, r이 너무 크므로 줄인다.
• 만약 합이 0보다 작다면, l이 너무 작으므로 늘린다.