[LeetCode] 1334. Find the City With the Smallest Number of Neighbors at a Threshold Distance (swift, python)

Find the City With the Smallest Number of Neighbors at a Threshold Distance - LeetCode

 

풀이

 각 도시에서 distanceThreshold 이하의 거리를 갖고 있는 도시의 개수를 카운트 해야 한다. 즉, '모든 지점에서 다른 모든 지점까지의 최단 경로를 모두 구해야 하는 경우'에 해당한다. 그렇기 때문에 최단 경로 알고리즘에서 플로이드 워셜 알고리즘을 사용해서 풀 수 있는 문제.

 각 도시에서 다른 모든 도시까지의 최단 경로를 구한 후 0~n-1번 도시 까지 distanceThreshold 이하의 거리로 갈 수 있는 도시를 각 도시 별로 찾아준다. 문제에서 가장 적은 개수를 가진 도시를 찾고 개수가 같은 도시가 있을 경우 가장 큰 도시를 출력하도록 요구하고 있으므로 플로이드 워셜 알고리즘을 적용한 후에 0번부터 for문을 통해서 적은 개수 이하의 도시가 있을 경우 적절히 갱신한 후 마지막에 구한 도시를 리턴해준다.

 

코드

[Swift]

class Solution {
    func findTheCity(_ n: Int, _ edges: [[Int]], _ distanceThreshold: Int) -> Int {
        let INF: Int = Int(1e9)
        var graph: [[Int]] = Array(repeating: Array(repeating: INF, count: n), count: n)
        
        // i에서 i로 가는 경우 0으로 설정
        for i in 0..<n {
            graph[i][i] = 0
        }
        
        // graph 배열에 weight 설정
        for edge in edges {
            graph[edge[0]][edge[1]] = edge[2]
            graph[edge[1]][edge[0]] = edge[2]
        }
        
        // 플로이드 워셜 알고리즘 적용
        for k in 0..<n {
            for i in 0..<n {
                for j in 0..<n {
                    graph[i][j] = min(graph[i][j], graph[i][k] + graph[k][j])
                }
            }
        }
        
        var minCnt: Int = INF
        var res: Int = 0
        
        // 각 노드 별로 distanceThreshold보다 적은 거리에 있는 노드의 갯수를 카운트해서 res 갱신
        for i in 0..<n {
            let cnt: Int = graph[i].filter() { $0 <= distanceThreshold }.count
            
            if cnt <= minCnt {
                minCnt = cnt
                res = i
            }
        }
        
        return res
    }
}
 

 

[Python]

 

INF = int(1e9)
 
 
class Solution:
    def findTheCity(self, n: int, edges: List[List[int]], distanceThreshold: int) -> int:
        graph = [[INF] * n for _ in range(n)]
 
        # i에서 i로 가는 경우는 0으로 설정
        for i in range(n):
            graph[i][i] = 0
 
        # graph 리스트에 weight 설정
        for (f, to, weight) in edges:
            graph[f][to] = weight
            graph[to][f] = weight
 
        # 플로이드 워셜 알고리즘 적용
        for k in range(n):
            for i in range(n):
                for j in range(n):
                    graph[i][j] = min(graph[i][j], graph[i][k] + graph[k][j])
 
        min_cnt = INF
        res = 0
        
        # 각 노드 별로 distanceThreshold보다 적은 거리에 있는 노드의 갯수를 카운트해서 res 갱신
        for i in range(n):
            cnt = len(list(filter(lambda x: x <= distanceThreshold, graph[i])))
 
            if cnt <= min_cnt:
                res = i
                min_cnt = cnt
 
        return res