viterbi.py

import sys

# 各ノードのコスト
nodes = {"BOS": 0, "EOS": 0,
         "さ": 100, "さか": 200, "さかな": 100, "かな": 200, "なだ": 200, "だ": 10, "よ": 10}

# ノードのつながりのコスト
edges = {("BOS", "さ"): 30,
         ("BOS", "さか"): 30,
         ("BOS", "さかな"): 30,
         ("さ", "かな"): 30,
         ("さか", "なだ"): 30,
         ("さかな", "だ"): 30,
         ("かな", "だ"): 30,
         ("なだ", "よ"): 45,
         ("だ", "よ"): 30,
         ("よ", "EOS"): 30}

# BOSから各nodeまでにかかるコスト
cost_dict = {node: sys.maxsize for node in nodes.keys()}
cost_dict["BOS"] = 0

# {node: 最短経路を通るときのnodeの前のノード}
route = {}


def get_prev_nodes(node):
    # nodeの前のnodeのリストを取得
    return [prev_node for prev_node, next_node in edges.keys() if next_node == node]


def main():
    # 最短経路の算出
    for node in ["さ", "さか", "さかな", "かな", "なだ", "だ", "よ", "EOS"]:
        shortest_prev_node = None
        min_cost = sys.maxsize

        # nodeの最小コストをもとめる
        for prev_node in get_prev_nodes(node):
            # startからprev_nodeまでのコスト + prev_nodeからnodeまでのエッジのコスト
            cost = cost_dict[prev_node] + edges[(prev_node, node)]
            if cost < min_cost:
                min_cost = cost
                shortest_prev_node = prev_node

        cost_dict[node] = min_cost + nodes[node]
        route[node] = shortest_prev_node

    # 最短経路の表示
    # BOSにたどりつくまでprev_nodeを追加していく
    shortest_path = ["EOS"]
    while shortest_path[0] != "BOS":
        prev_node = route[shortest_path[0]]
        shortest_path.insert(0, prev_node)
    print("shortest path:", shortest_path)


if __name__ == "__main__":
    main()