백준 7453번 합이 0인 네 정수 v2

문제 링크: https://www.acmicpc.net/problem/7453

백준 7453번 합이 0인 네 정수 v2/main.py

@@ -0,0 +1,50 @@
+import sys
+
+
+def solve():
+    input = sys.stdin.readline
+    N = int(input().strip())
+    dts = [list(map(int, input().strip().split())) for _ in range(N)]
+    ABs, CDs = [], []
+
+    for i in range(N):
+        for j in range(N):
+            ABs.append(dts[i][0] + dts[j][1])
+            CDs.append(dts[i][2] + dts[j][3])
+
+    ABs.sort()
+    CDs.sort()
+
+    idx_ab = 0
+    idx_cd = (N * N) - 1
+    answer = 0
+
+    while idx_ab < (N * N) and 0 <= idx_cd:
+        v = ABs[idx_ab] + CDs[idx_cd]
+
+        if v == 0:
+            idx_ab2 = idx_ab + 1
+            while idx_ab2 < (N * N) and ABs[idx_ab] == ABs[idx_ab2]:
+                idx_ab2 += 1
+            len_ab = idx_ab2 - idx_ab
+
+            idx_cd2 = idx_cd - 1
+            while 0 <= idx_cd2 and CDs[idx_cd] == CDs[idx_cd2]:
+                idx_cd2 -= 1
+            len_cd = idx_cd - idx_cd2
+
+            answer += (len_ab * len_cd)
+
+            idx_ab = idx_ab2
+            idx_cd = idx_cd2
+
+        elif v < 0:
+            idx_ab += 1
+        else:  # V > 0
+            idx_cd -= 1
+
+    print(answer)
+
+
+if __name__ == '__main__':
+    solve()

백준 7453번 합이 0인 네 정수 v2/test1.txt

@@ -0,0 +1,7 @@
+6
+-45 22 42 -16
+-41 -27 56 30
+-36 53 -37 77
+-36 30 -75 -46
+26 -38 -10 62
+-32 -54 -6 45

백준 7453번 합이 0인 네 정수 v2/test1_answer.txt

@@ -0,0 +1 @@
+5

백준 7453번 합이 0인 네 정수 v2/test_main.py

@@ -0,0 +1,23 @@
+import sys
+from pathlib import Path
+from unittest import TestCase
+from main import solve
+
+
+class Test(TestCase):
+    def my_solve(self, testcase_input):
+        sys.stdin = open(testcase_input, 'r')
+        stdout = sys.stdout
+        sys.stdout = open('stdout.txt', 'w')
+        solve()
+        sys.stdout.close()
+        sys.stdout = stdout
+
+    def test_solve(self, testcase_number: str):
+        self.my_solve('test' + testcase_number + '.txt')
+        self.assertEqual(
+            Path('test' + testcase_number + '_answer.txt').read_text().strip(),
+            Path('stdout.txt').read_text().strip())
+
+    def test1_solve(self):
+        self.test_solve('1')