[Python] 소프티어 : 거리 합 구하기

https://softeer.ai/practice/info.do?idx=1&eid=635&sw_prbl_sbms_sn=157103 

Softeer

---

💡 문제 풀이

그래프 탐색 문제이다.

해결하지 못했다.

✔️ 느낀 점

보자마자 플로이드-워샬이 생각나서 접근했지만 시간초과로 실패 (O(n^3)),

DFS로 접근했는데 중간 계산 로직을 구현하지 못해서 실패,

다익스트라로 접근했는데 개선된 버전을 사용해도 시간초과로 실패 (O(E^2logE))

DFS가 두 번 들어가면 풀린다고 하는데 솔직히 좋은 문제인지는 잘 모르겠다.

💻 코드

import sys
input = sys.stdin.readline

INF = int(1e9)
n = int(input())
graph = [[INF] * (n+1) for _ in range(n+1)]
for i in range(1, n+1):
    graph[i][i] = 0

for _ in range(n-1):
    x, y, c = map(int, input().split())
    graph[x][y] = c
    graph[y][x] = c
    
for k in range(1, n+1):
    for a in range(1, n+1):
        for b in range(1, n+1):
            graph[a][b] = min(graph[a][b], graph[a][k] + graph[k][b])
            
for i in range(1, n+1):
    sum_ = 0
    for j in range(1, n+1):
        if graph[i][j] == INF: continue
        sum_ += graph[i][j]
    print(sum_)

import sys
input = sys.stdin.readline

def dfs(i, graph):
    totalCost, particalCost = 0, 0
    stack, visited = [], []
    stack.append([i, 0])
    
    while stack:
        node, cost = stack.pop()
        
        if node not in visited:
            stack.extend(graph[node])
            visited.append(node)
            particalCost += cost
            totalCost += particalCost
        else:
            particalCost -= cost
        
    return totalCost
        
n = int(input())
graph = [[] for _ in range(n+1)]

for _ in range(n-1):
    s, e, c = map(int, input().split())
    graph[s].append([e,c])
    graph[e].append([s,c])
    
for i in range(1, n+1):
    print(dfs(i, graph))

import sys, heapq
input = sys.stdin.readline

INF = int(1e9)

def dijkstra(graph, start, size):
    distance = [INF] * (size+1)
    distance[start] = 0
    queue = []
    
    heapq.heappush(queue, [0, start])
    
    while queue:
        currentCost, currentNode = heapq.heappop(queue)
        
        if distance[currentNode] < currentCost: continue
        
        for nextCost, nextNode in graph[currentNode]:
            newCost = nextCost + currentCost
            if newCost < distance[nextNode]:
                distance[nextNode] = newCost
                heapq.heappush(queue, [newCost, nextNode])
    
    sum_ = 0
    for cost in distance:
        if cost == INF: continue
        sum_ += cost
        
    return sum_

n = int(input())
graph = [[] for _ in range(n+1)]

for _ in range(n-1):
    start, end, cost = map(int, input().split())
    graph[start].append([cost, end])
    graph[end].append([cost, start])
    
for i in range(1, n+1):
    print(dijkstra(graph, i, n))

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**6)

def dfs1(current, parent, subtree_size, dist_sum):
    subtree_size[current] = 1
    for i in range(len(graph[current])):
        child = graph[current][i][0]
        weight = graph[current][i][1]
        
        if child != parent:
            dfs1(child, current, subtree_size, dist_sum)
            subtree_size[current] += subtree_size[child]
            dist_sum[current] += dist_sum[child] + weight * subtree_size[child]
            
    return subtree_size,dist_sum

def dfs2(current, parent, subtree_size, dist_sum):
    for i in range(len(graph[current])):
        child = graph[current][i][0]
        weight = graph[current][i][1]
        
        if child != parent:
            dist_sum[child] = dist_sum[current] + weight * (n - subtree_size[child]) - weight * subtree_size[child]
            dfs2(child, current, subtree_size, dist_sum)
            
    return dist_sum

n = int(input())
graph = [[] for _ in range(n+1)]

for _ in range(n-1):
    x,y,t = map(int,input().split())
    graph[x].append((y,t))
    graph[y].append((x,t))
    
subtree_size, dist_sum = dfs1(1, 1, [0]*(n+1), [0]*(n+1))
dist_sum = dfs2(1, 1, subtree_size, dist_sum)

for i in range(1,n+1):
    print(dist_sum[i])