[Python/파이썬] 백준 1711번 직각삼각형

1711번: 직각삼각형

 

문제

2차원 평면에 N개의 점이 주어져 있다. 이 중에서 세 점을 골랐을 때, 직각삼각형이 몇 개나 있는지를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 점의 개수 N(3 ≤ N ≤ 1,500)이 주어진다. 둘째 줄부터 N개의 줄에 걸쳐 각 점의 x좌표와 y좌표가 빈 칸을 사이에 두고 주어진다. 좌표값은 절댓값이 1,000,000,000을 넘지 않는 정수이며, 주어지는 모든 점의 좌표는 다르다.

출력

첫째 줄에 직각삼각형의 개수를 출력한다.

코드

n = int(input())
dots = []
for _ in range(n):
    x, y = map(int, input().split())
    dots.append((x, y))

# 선분 ab의 길이 ^ 2
def length(a, b):
    return (a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2

# 직각삼각형인지
def chk(a, b, c):
    ab = length(a, b)
    bc = length(b, c)
    ca = length(a, c)
    # a가 직각
    if ab + ca == bc:
        return True
    # b가 직각
    elif ab + bc == ca:
        return True
    # c가 직각
    elif ca + bc == ab:
        return True
    return False
    
ans = 0
for i in range(n-2):
    for j in range(i+1, n-1):
        for k in range(j+1, n):
            if chk(dots[i], dots[j], dots[k]):
                ans += 1
print(ans)

Python3로는 시간 초과 발생

PyPy3로 통과 가능

 

처음에는 itertools 라이브러리의 combinations를 이용하여 dots에서 3개를 뽑는 조합을 만들어서 했는데 Python3로는 물론이고 PyPy3에서도 시간 초과가 발생했다. 그래서 3중 for문으로 변경했더니 PyPy3에서는 간신히 통과...