[백준/Python] Silver I #20916 안녕 2020 안녕 2021

20916번: 안녕 2020 안녕 2021

문제

2020년을 보내고 2021년을 맞이하는 기념으로 Albert는 재미있는 문제를 풀기로 했다. 양의 정수 중 첫 네 숫자가 "2020"이며 마지막 네 숫자가 "2021"이면 "안녕한 정수" 라고 정의 하자. 가령 202021 이나 20202021은 "안녕한 정수" 이며, 2020021 이나 2020221 은 안녕한 정수가 아니다.

Albert는 n개의 정수 A[1], A[2], ..., A[n]중 두 수를 골라 더했을 때 그 합이 "안녕한 정수"가 되는 쌍의 개수를 알고 싶다. 즉, 1 ≤ i < j ≤ n 을 만족하는 (i, j) 쌍 중 A[i] + A[j]가 안녕한 정수인 쌍의 개수를 알고 싶다. 예를 들어 A = [101010, 101010, 101011, 101011], 즉 n = 4개의 정수가 있다고 하자. 이 경우 A[1] + A[3] = A[1] + A[4] = A[2] + A[3] = A[2] + A[4] = 202021 이므로 총 4개의 쌍이 존재한다 ((1, 3), (1, 4), (2, 3), (2, 4)).

입력으로 n개의 정수가 주어졌을 때, 합이 안녕한 정수가 되도록 하는 쌍의 개수를 출력하시오.

입력

첫 줄에 테스트 케이스의 수 T가 주어진다. 각 테스트 케이스는 두 줄로 구성된다.

첫 줄에 정수의 개수 n이 주어진다. 다음 줄에 n개의 정수가 공백으로 구분되어 주어진다.

출력

각 테스트 케이스에 대해 합이 안녕한 정수가 되는 쌍의 개수를 출력한다.

풀이

target_numbers = [202021, 20202021, 202002021, 202012021, 202022021, 202032021, 202042021, 202052021, 202062021, 202072021, 202082021, 202092021]

T = int(input())
for _ in range(T):
    n = int(input())
    arr = list(map(int, input().split()))
    arr.sort()

    freq = {}
    for num in arr:
        freq[num] = freq.get(num, 0) + 1

    ans = 0
    for num in arr:
        for target in target_numbers:
            complement = target - num
            if complement in freq:
                ans += freq[complement]
                if complement == num:
                    ans -= 1
    print(ans // 2)

시간초과 때문에 꽤나 힘들었던 문제... 시간초과를 해결하려고 하면 메모리초과가 뜨고.. 무한 굴레에 빠졌다.
심지어 맞은 사람도 적고, 풀이를 올린 사람은 단 한명도 없어서 깃헙에 누군가가 올렸을거란 생각으로 엄청 찾다가 구세주를 찾게 되었다.
하지만 그 분은 C++로 풀어서 파이썬으로 바꾸는 과정이 필요했고, 이 과정에서 lower_bound와 upper_bound를 알게 되었다.
이해해서 파이썬으로 바꿔서 제출했는데도 시간초과가 나서..... 시간을 줄일만한 방법을 찾아서 풀었더니 맞았습니다! ..
진짜 힘들었다...... 20000번대 문제는.. 잘 골라서 풀어야겠다.