백준 / 텔레포트 / 16958번 / Python

*문제 출처는 백준에 있습니다.

문제 제목: 텔레포트 / 16958번 (골드 4단계)

문제 사이트: https://www.acmicpc.net/problem/16958

 

문제 설명

---

나의 풀이

import sys
input = sys.stdin.readline

INF = 10**15

def manhattan(x1, y1, x2, y2):
    return abs(x1 - x2) + abs(y1 - y2)

if __name__ == "__main__":
    n, t = map(int, input().split())

    # 1-indexed 저장: (s, x, y)
    cities = [None]
    specials = []  # (x, y)만 저장
    for _ in range(n):
        s, x, y = map(int, input().split())
        cities.append((s, x, y))
        if s == 1:
            specials.append((x, y))

    # 각 도시에서 "가장 가까운 특별도시"까지의 거리 전처리
    # 특별도시가 없으면 텔레포트 경로는 사용할 수 없음(INF로 유지)
    nearest_special = [INF] * (n + 1)
    if specials:
        for i in range(1, n + 1):
            s, x, y = cities[i]
            if s == 1:
                nearest_special[i] = 0
            else:
                best = INF
                for sx, sy in specials:
                    d = abs(x - sx) + abs(y - sy)
                    if d < best:
                        best = d
                nearest_special[i] = best

    m = int(input())
    out = []
    for _ in range(m):
        a, b = map(int, input().split())
        _, x1, y1 = cities[a]
        _, x2, y2 = cities[b]

        direct = manhattan(x1, y1, x2, y2)
        via_tp = nearest_special[a] + t + nearest_special[b]  # specials 없으면 INF라 자동 배제
        out.append(str(min(direct, via_tp)))

    print("\n".join(out))

---

※ 알아야 할 것

import sys
INF = int(1e9)  # 10억


def prim(citys):
    distance = [[INF] * (len(citys) + 1) for _ in range(len(citys) + 1)]  # 거리

    # 거리 갱신
    for i in range(1, n + 1):
        for j in range(1, n + 1):
            if i == j:
                distance[i][j] = 0
            s1, x1, y1 = citys[i]
            s2, x2, y2 = citys[j]

            direction = abs(x2 - x1) + abs(y2 - y1)

            if s1 == 1 and s2 == 1:
                tp = t
                min_value = min(direction, tp)
                distance[i][j] = min(distance[i][j], min_value)
            else:
                distance[i][j] = min(distance[i][j], direction)

    # 자기자신은 거리 0
    for k in range(1, n + 1):
        for i in range(1, n + 1):
            for j in range(1, n + 1):
                distance[i][j] = min(distance[i][j], distance[i][k] + distance[k][j])

    return distance


if __name__=="__main__":
    # 도시의 수 n, 텔레포트하는데 걸리는 시간 t
    n, t = map(int, input().split())

    # 1번 도시부터 시작
    city = [[]]

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

    dis = prim(city)

    m = int(input())
    result = []

    for _ in range(m):
        a, b = map(int, input().split())
        result.append(dis[a][b])

    for i in result:
        print(i)

위 코드는 실제로 플로이드 워셜 알고리즘을 이용해서 구현을 했지만 O(10^9) 시간복잡도가 나와서 틀렸습니다.