tgraph.py

import collections
import heapq

from typing import Mapping, Sequence, Optional, Iterable, List

Graph = Sequence[Iterable[int]]
WGraph = Sequence[Mapping[int, float]]

INF = float('inf')

# N dfs
# I {"version": "1.0", "typing": ["List", "Generator", "Iterable", "Optional", "Sequence"], "const": ["Graph"]}
def dfs(graph: Graph,
        source: int,
        stack: Optional[List] = None,
        yields_on_leave: bool = False) -> Generator[int, None, None]:    
    is_visited = [False] * len(graph)            
    if stack is None:
        stack = []
    is_visited[source] = True            
    stack.append(source)        
    iter_stack = [iter(graph[source])]
    yield source
    while iter_stack:
        try:
            node = next(iter_stack[-1])
            if not is_visited[node]:
                is_visited[node] = True
                stack.append(node)
                iter_stack.append(iter(graph[node]))
                yield node
        except StopIteration:
            if yields_on_leave:
                yield stack[-1]
            stack.pop()
            iter_stack.pop()

• non-recursive DFS

# N min_distances
# I {"version": "1.0", "import": ["collections"], "typing": ["List", "Iterable", "Optional", "Sequence"], "const": ["Graph", "INF"]}
def min_distances(graph: Graph,
                  source: int,
                  sink: Optional[int] = None) -> List[int]:
    """Returns shortest path distances from source node to all nodes."""
    distances = [INF] * len(graph)
    distances[source] = 0
    queue = collections.deque([(source, 0)])
    while queue:
        u, dist = queue.popleft()
        for v in graph[u]:
            if distances[v] == INF:
                distances[v] = dist + 1
                if v == sink:
                    return distances
                queue.append((v, dist + 1))
    return distances

• 단일 출발지 최단 경로 (Single Source Shortest Path) 참고

# N min_distances_on_tree
# I {"version": "1.0", "typing": ["List", "Mapping", "Optional", "Sequence"], "const": ["WGraph", "INF"]}
def min_distances_on_tree(wgraph: WGraph,
                          source: int,
                          sink: Optional[int] = None) -> List[float]:
    """Returns shortest path distances from source node to all nodes."""
    distances = [INF] * len(wgraph)
    distances[source] = 0
    stack = [(source, 0, None)]
    while stack:
        u, dist_to_u, parent = stack.pop()
        distances[u] = dist_to_u
        if u == sink:
            return distances
        for v, dist in wgraph[u].items():
            if v != parent:
                stack.append((v, dist_to_u + dist, u))
    return distances

• 단일 출발지 최단 경로 (Single Source Shortest Path) 참고

# N minimum_spanning_tree
# I {"version": "1.0", "import": ["heapq"], "typing": ["Mapping", "Sequence"], "const": ["WGraph"]}
def minimum_spanning_tree(wgraph: WGraph) -> int:
    """Returns sum of weights in the minimum spanning tree of given graph.

    Based on Prim's algorithm with binary heap, running in O(ElogV)."""
    total_cost = 0
    is_visited = [False] * len(wgraph)
    heap = [(0, 0)]
    while heap:
        cost, u = heapq.heappop(heap)
        if is_visited[u]:
            continue
        is_visited[u] = True
        total_cost += cost
        for v, cost_to_v in wgraph[u].items():
            heapq.heappush(heap, (cost_to_v, v))
    return total_cost

• 최소 신장 트리 (Minimum Spanning Tree / MST) 참고

# N all_pairs_shortest_path
# I {"version": "1.0", "typing": ["List", "Mapping", "Sequence"], "const": ["WGraph", "INF"]}
def all_pairs_shortest_path(wgraph: WGraph) -> List[List[int]]:
    """Returns a table that contains shortest distances between all nodes."""
    n = len(wgraph)
    dists = [[INF] * n for _ in range(n)]
    for i in range(n):
        dists[i][i] = 0
    for dists_u, edges in zip(dists, wgraph):
        for v, dist_to_v in edges.items():
            dists_u[v] = dist_to_v

    for k, dists_k in enumerate(dists):
        for dists_i in dists:
            dists_ik = dists_i[k]            
            for j, (dists_ij, dists_kj) in enumerate(zip(dists_i, dists_k)):
                if dists_ik + dists_kj < dists_ij:
                    dists_i[j] = dists_ik + dists_kj
    return dists

• 전체쌍 최단 경로 (All-pairs shortest path) 참고
• 인풋으로 인접행렬 형태의 그래프를 받으면, 지금처럼 인접리스트를 인접행렬로 변환하는 과정이 필요없으므로 더 빠르겠지만, 다른 함수들과의 일관성을 위해서 인접리스트를 받는 것으로 했다.
• 대신 속도를 만회하기 위해서 코드를 좀더 최적화 했다. O(n^3)이다보니 사소한 최적화에도 실행속도가 꽤 영향을 받는다.

# N dijkstra
# I {"version": "1.0", "import": ["heapq"], "typing": ["List", "Mapping", "Optional", "Sequence"], "const": ["WGraph", "INF"]}
def dijkstra(wgraph: WGraph,
             source: int,
             sink: Optional[int] = None) -> List[float]:
    """Returns a list of minimum costs from source to each node.

    This is an implementation of Dijkstra algorithm using binary heap, running
    in O(ElogV).
    """
    distances = [INF] * len(wgraph)
    distances[source] = 0
    heap = [(0, source)]
    while heap:
        dist_u, u = heapq.heappop(heap)
        if dist_u != distances[u]:
            continue
        if u == sink:
            return distances
        for v, dist_uv in wgraph[u].items():
            dist_v = dist_u + dist_uv
            if dist_v < distances[v]:
                distances[v] = dist_v
                heapq.heappush(heap, (dist_v, v))
    return distances

• 다익스트라 알고리즘 참고

# N spfa
# I {"version": "1.4", "import": ["collections"], "typing": ["List", "Mapping", "Optional", "Sequence"], "const": ["WGraph", "INF"]}
def spfa(wgraph: WGraph,
         source: int,
         prev: Optional[List[int]] = None) -> List[float]:
    size = len(wgraph)
    lengths = [0] * size
    is_in_queue = [False] * size
    is_in_queue[source] = True
    distances = [INF] * size
    distances[source] = 0
    if prev:
        prev[source] = None
    queue = collections.deque([source])
    while queue:
        u = queue.popleft()
        is_in_queue[u] = False
        length_u, dist_u = lengths[u], distances[u]
        if length_u == size:
            dist_u = distances[u] = -INF
        for v, dist_uv in wgraph[u].items():
            dist_v = dist_u + dist_uv
            if distances[v] <= dist_v:
                continue
            distances[v], lengths[v] = dist_v, length_u + 1
            if prev:
                prev[v] = u
            if not is_in_queue[v]:
                queue.append(v)
                is_in_queue[v] = True
    return distances

• Shortest Path Faster Algorithm (SPFA) 참고