백준: 7453 합이 0인 네 정수

문제

7453번: 합이 0인 네 정수

 

문제 풀이

 앞 서 다룬 부분수열의 합 2, 두 배열의 합 합과 동일한 방식으로 푸는 문제이다. 한 줄에 A, B, C, D 4 쌍의 값이 있고 N개의 줄에 걸쳐 여러 쌍의 정수들이 주어진다. 앞의 문제들과 마찬가지로 한 번에 모든 경우를 구하는 것은 시간 초과가 발생한다. 따라서 A, B에 대한 부분합에 대한 카운트를 진행한 후, 0 - C, D의 부분합의 경우의 수를 정답으로 반영하면 문제를 풀 수 있다.

 

코드

from sys import stdin
from collections import defaultdict


def subtotal(lst, select):
    global answer

    for i in lst[0]:
        for j in lst[1]:
            if select == "ab":
                subtotal_a_b[i + j] += 1
            else:
                answer += subtotal_a_b[0 - (i + j)]


if __name__ == "__main__":
    answer = 0
    subtotal_a_b = defaultdict(int)
    n = int(stdin.readline())
    nums = [list(map(int, stdin.readline().split())) for _ in range(n)]
    nums = list(map(list, zip(*nums)))
    subtotal(nums[:2], "ab")
    subtotal(nums[2:], "cd")
    print(answer)

 입력된 값을 [A, B], [C, D]로 분류하기 위해 행열을 전치시켰다. Numpy 없이 행열을 전치시키기 위해서는 `map`과 `zip`을 활용하면 `list(map(list, zip(*전치할 행열)))`과 같이 한 줄로 전치시킬 수 있다.