카카오 2021 기출 - 합승 택시요금 [워셜 플로이드]

def solution(numOfPoint, S, A, B, edges) :
    costs = [[0 if i==j else 1000000000 for i in range(numOfPoint)] for j in range(numOfPoint)]
    for edge in edges:
        costs[edge[0]-1][edge[1]-1] = edge[2]
        costs[edge[1]-1][edge[0]-1] = edge[2]
    for middlePoint in range(numOfPoint):
        for start in range(numOfPoint):
            for end in range(numOfPoint):
                if middlePoint != start and middlePoint != end and start != end:
                    costs[start][end] = min(costs[start][end], costs[start][middlePoint] + costs[middlePoint][end])
    answer = 10000000
    for separatePoint in range(numOfPoint):
        cost = costs[S-1][separatePoint] + costs[separatePoint][A-1] + costs[separatePoint][B-1]
        answer = answer if answer < cost else cost
    return answer



print(solution(6, 4, 6, 2, [[4, 1, 10], [3, 5, 24], [5, 6, 2], [3, 1, 41], [5, 1, 24], [4, 6, 50], [2, 4, 66], [2, 3, 22], [1, 6, 25]]))

워셜 플로이드는 시간 복잡도 n^3 이다 그래서 그런가

26번이 시간초과..

그래서 다음과 같이 다익스트라로 바꿔봤다.

import heapq

INF = 1000000000

def dijkstra(start, costs, distance):
    q = []
    heapq.heappush(q, (0, start))
    distance[start][start] = 0
    while q:
        dist, cur = heapq.heappop(q)
        if distance[start][cur] < dist:
            continue
        for idx, i in enumerate(costs[cur]):
            throughDist = dist + i
            if distance[start][idx] > throughDist :
                distance[start][idx] = throughDist
                heapq.heappush(q, (throughDist,idx))


def solution(numOfPoint, S, A, B, edges) :
    costs = [[0 if i==j else INF for i in range(numOfPoint)] for j in range(numOfPoint)]
    for edge in edges:
        costs[edge[0]-1][edge[1]-1] = edge[2]
        costs[edge[1]-1][edge[0]-1] = edge[2]

    distance = [[0 if i==j else INF for i in range(numOfPoint)] for j in range(numOfPoint)]
    for start in range(numOfPoint):
        dijkstra(start, costs, distance)

    answer = 10000000
    for separatePoint in range(numOfPoint):
        cost = distance[S-1][separatePoint] + distance[separatePoint][A-1] + distance[separatePoint][B-1]
        answer = answer if answer < cost else cost
    return answer



print(solution(6, 4, 6, 2, [[4, 1, 10], [3, 5, 24], [5, 6, 2], [3, 1, 41], [5, 1, 24], [4, 6, 50], [2, 4, 66], [2, 3, 22], [1, 6, 25]]))

성공!

 

근데 다 풀고 나서 워셜 플로이드 알고리즘에 min 함수를 안쓰고 해봤더니 통과...... 충격

 

오늘의 값진 결론

단순 값 2개 비교시에는 min, max함수를 쓰지말자. 덕분에 2시간 뻘짓한거 보상받은 기분이다

 

def solution(numOfPoint, S, A, B, edges) :
    costs = [[0 if i==j else 1000000000 for i in range(numOfPoint)] for j in range(numOfPoint)]
    for edge in edges:
        costs[edge[0]-1][edge[1]-1] = edge[2]
        costs[edge[1]-1][edge[0]-1] = edge[2]
    for middlePoint in range(numOfPoint):
        for start in range(numOfPoint):
            for end in range(numOfPoint):
                if middlePoint != start and middlePoint != end and start != end:
                    midCost = costs[start][middlePoint] + costs[middlePoint][end]
                    costs[start][end] = costs[start][end] if midCost > costs[start][end] else midCost
    answer = 10000000
    for separatePoint in range(numOfPoint):
        cost = costs[S-1][separatePoint] + costs[separatePoint][A-1] + costs[separatePoint][B-1]
        answer = answer if answer < cost else cost
    return answer

최종 코드