[실버 3] 백준 3273 - 두 수의 합 (파이썬)

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

풀이

• 정수 수열 A에서 서로 다른 두 수의 합이 X가 되는 경우의 수를 구하는 문제이다.
• 수의 개수 N은 최대 10만개
• 수의 범위는 1이상 1백만 이하

진짜 단순하게 itertools의 combination을 사용하면 시간 초과가 엄청나게 난다는 사실,,,

# 백준 3273 - 두수의 합
# 분류 - 투포인터

import sys
from itertools import combinations
input = sys.stdin.readline

N = int(input())
A = list(map(int, input().split()))
X = int(input())
count = 0
for combination in combinations(A, 2) :
    if (sum(combination)) == X :
        count += 1
print(count)

처음 짠 코드는 단순히 combinations를 사용하여 X와 같을 경우는 count += 1을 진행해주었지만, 시간이 1초로 제한되어 있기때문에 당연히 시간초과가 일어났다.

그래서 아래와 같이 수정을 진행했다.

count = [0] * 1000001
for i in range(N):
    count[A[i]] += 1

이 코드는 각 수의 등장 횟수를 세기 위한 코드이다.

숫자들이 1부터 1백만 사이이므로, 등장 횟수를 배열로 저장한다.

for i in range((X + 1) // 2):
    if X - i >= 1000000:
        continue
    answer += count[i] * count[X - i]

서로 다른 두 수의 합이 X가 되는 쌍의 수를 구하기 위해서 다음과 같은 조건을 고려했다.

1. i + (X - i) = X 이므로, i와 X - i의 조합을 고려한다.
2. (X + 1) // 2 까지만 반복해서 중복 카운팅을 방지한다.
3. count[i] * count[X - i]는 가능한 쌍의 개수를 처리한다.

if X % 2 == 0:
    answer += count[X // 2] * (count[X // 2] - 1) // 2

X가 짝수일 때 중앙 값 하나로 이루어진 조합을 처리하기 위해서 사용한 코드지만, 사실상 이 문제에서는 모든 정수가 서로 다르다고 하였기 때문에 큰 의미는 없다.

시간 복잡도 분석

• 입력 : O(N)
• 카운팅 탐색 : O(X)

코드

# 백준 3273 - 두수의 합
# 분류 - 투포인터

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

count = [0] * 1000001
for i in range(N) :
    count[A[i]] += 1

answer = 0
for i in range((X + 1) // 2) :
    if X - i >= 1000000 :
        continue
    answer += count[i] * count[X - i]

if X % 2 == 0 :
    answer += count[X // 2] * (count[X // 2] - 1) // 2

print(answer)