[백준: Python] 7453 - 합이 0인 네 정수

[문제]

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

시간제한이 무려 12초인 문제입니다.

 

대놓고 최적화를 요구하는 문제라는 것을 알 수 있습니다.

 

특정 값을 만들 수 있는 조합의 수를 찾는 것이 목표입니다.

 

이러한 문제 유형의 경우, 전부 계산하지 않고 부분으로 쪼개 해시에 저장해 값을 비교하는 것이 일반적입니다.

 

부르트포스하게 풀면 O(n^4)이지만 해시를 사용하면 O(n^2)에 풀 수 있습니다.

---

[초기화]

    def __init__(self):
        self.n = int(input())
        self.arr = [list(map(int, input().split())) for _ in range(self.n)]
        self.a, self.b, self.c, self.d = [], [], [], []
        self.hash = dict()
        self.answer = 0

    def init_arr(self):
        for i in self.arr:
            self.a.append(i[0])
            self.b.append(i[1])
            self.c.append(i[2])
            self.d.append(i[3])

처음에는 입력예제가 바로 이해되지 않았는데,  입력들을 하나의 열로 생각했을 때 순서대로 A, B, C, D가 됩니다.

---

[풀이]

    def solve(self):
        self.init_arr()

        for a in self.a:
            for b in self.b:
                key = a+b
                if key in self.hash:  # a+b를 키를 가지는 해시맵 생성
                    self.hash[key] += 1
                else:
                    self.hash[key] = 1

        for c in self.c:
            for d in self.d:
                key = -(c+d)  # c+d를 음수를 취한 값이 존재하는지 확인, -(c+d)가 존재한다는 것은 a+b가 c+d와 같지만 반대인 것을 의미
                if key in self.hash:
                    self.answer += self.hash[key]

        print(self.answer)

먼저 a + b의 조합 쌍을 찾으며 해시에 기록합니다.

 

다음 c + d의 조합 쌍을 찾고, 이것에 대응되는(a + b = -(c + d)) 쌍의 수를 합산합니다.

 

위와 같이 되는 이유는 a + b + c + d = 0이 되는 것만 찾으면 되기 때문입니다. 

---

[전체 코드]

import sys
input = lambda: sys.stdin.readline().rstrip()


class Main:
    def __init__(self):
        self.n = int(input())
        self.arr = [list(map(int, input().split())) for _ in range(self.n)]
        self.a, self.b, self.c, self.d = [], [], [], []
        self.hash = dict()
        self.answer = 0

    def init_arr(self):
        for i in self.arr:
            self.a.append(i[0])
            self.b.append(i[1])
            self.c.append(i[2])
            self.d.append(i[3])

    def solve(self):
        self.init_arr()

        for a in self.a:
            for b in self.b:
                key = a+b
                if key in self.hash:  # a+b를 키를 가지는 해시맵 생성
                    self.hash[key] += 1
                else:
                    self.hash[key] = 1

        for c in self.c:
            for d in self.d:
                key = -(c+d)  # c+d를 음수를 취한 값이 존재하는지 확인, -(c+d)가 존재한다는 것은 a+b가 c+d와 같지만 반대인 것을 의미
                if key in self.hash:
                    self.answer += self.hash[key]

        print(self.answer)


problem = Main()
problem.solve()

---