Dijkstras for Maya

dijkstras_for_Maya.py

# -*- coding: utf-8 -*- 
#ダイクストラ法
from maya import cmds
from collections import OrderedDict
import maya.api.OpenMaya as om2
import math
import copy
import time

def dijkstras():
    #開始時間
    print 'Dijkstras algorithm------------------------------------'
    start = time.time()
    rap_time = time.time()
    
    #メッシュと探索頂点の開始終了位置を取得
    try:
        target_mesh = cmds.ls(hl=True)[0]
        start_vtx = cmds.ls(sl=True, fl=True)[0]
        goal_vtx = cmds.ls(sl=True, fl=True)[1]
    except:
        print '頂点を2つ選択してください'
        return
    
    #すべての頂点の位置を取得
    selection = om2.MGlobal.getActiveSelectionList()
    mDag, __ = selection.getComponent(0)
    if mDag.hasFn(om2.MFn.kMesh):
        mfnMesh = om2.MFnMesh(mDag)
    vertex_array = mfnMesh.getPoints(om2.MSpace.kWorld)
    
    #メッシュの全頂点を取得しておく
    all_vtx = cmds.polyListComponentConversion(target_mesh, tv=True)
    all_vtx = cmds.filterExpand(all_vtx, sm=31)
    
    #頂点位置をすべて格納しておく
    vtx_pos_dict = {vtx:pos for vtx, pos in zip(all_vtx, vertex_array)}
    
    print u'rap time / メッシュ頂点取得:', time.time() - rap_time
    rap_time = time.time()
    
    #全頂点のつながり辞書を作っておく、頂点をエッジ→頂点変換すると周辺の頂点が取れるので簡易取得。
    temp_Dict = {}
    for vtx in all_vtx:
        around = cmds.polyListComponentConversion(vtx, tf=True)
        around = cmds.polyListComponentConversion(around, tv=True)
        around = cmds.filterExpand(around, sm=31)
        around.remove(vtx)
        roundList = []
        a = vtx_pos_dict[vtx]
        for rVtx in around:
            b = vtx_pos_dict[rVtx]
            dist = (a-b).length()
            roundList.append((rVtx, dist))
        temp_Dict[vtx] = roundList
        
    print u'rap time / つながり辞書   :', time.time() - rap_time
    rap_time = time.time()
    
    #スタートとゴールの経路を保持するためオーダードディクトを利用
    route_dict = OrderedDict()
    #スタートとゴールを最初と最後のキーと要素に設定する。あとは適当。
    route_dict[start_vtx] = temp_Dict[start_vtx]
    for k in temp_Dict.keys():
        if k != start_vtx and k != goal_vtx:
            route_dict[k] = temp_Dict[k]
    route_dict[goal_vtx] = temp_Dict[goal_vtx]
    
    print 'rap time / ルート格納辞書 :', time.time() - rap_time
    rap_time = time.time()
    
    #探索用変数を色々初期化
    node_num = len(route_dict)#ノードの総数
    distance_dict = {vtx:float('inf') for vtx in all_vtx}#ノードごとの距離のリスト、初期値無限大を与える
    previous_nodes_dict = {vtx:'' for vtx in all_vtx}#最短経路でそのノードのひとつ前に到達するノードの対応辞書
    distance_dict[start_vtx] = 0#初期のノードの距離は0とする
    unsearched_nodes = copy.copy(distance_dict)#未探索ノード
    
    print 'rap time / データ初期化   :', time.time() - rap_time
    rap_time = time.time()
    
    i = 0
    while(len(unsearched_nodes) != 0): #未探索ノードがなくなるまで繰り返す
        #未探索ノードのうちで最小のindex番号を取得
        min_dist = min(unsearched_nodes.values())
        min_id = unsearched_nodes.values().index(min_dist)
        target_min_key = unsearched_nodes.keys()[min_id]
               
        #未探索ノードから除去しておく
        unsearched_nodes.pop(target_min_key)
        
        #ターゲットノードの周辺ノード取得
        roundVertex = route_dict[target_min_key]
        
        #周辺ノードへのスタートからの到達距離を比較、以前設定した距離より小さい場合は更新する。
        for rVtx in roundVertex:
            i += 1
            real_dist = distance_dict[target_min_key] + rVtx[1]#実際の距離
            if distance_dict[rVtx[0]] > real_dist:
                distance_dict[rVtx[0]] = real_dist
                unsearched_nodes[rVtx[0]] = real_dist
                #そのノードに最短で到達するための一つまえのノード辞書を更新しておく
                previous_nodes_dict[rVtx[0]] =  target_min_key
        
    print u'rap time / ルート探索     :', time.time() - rap_time
    rap_time = time.time()
                
    #最短ひとつ前辞書をゴールから逆引きしてスタートまでつなげる。
    minimum_route = []
    previous_node = goal_vtx
    while previous_node != start_vtx:
        minimum_route.append(previous_node)
        previous_node = previous_nodes_dict[previous_node]
        
    #ゴールへの最短到達距離を取得
    minimum_route.append(previous_node)
    
    #探索したつながりを選択
    cmds.select(minimum_route, r=True)
    elapsed_time = time.time() - start
    
    print u'rap time / 最終データ取得 :', time.time() - rap_time
    rap_time = time.time()
    
    #プリントして探索結果を確認
    print u'総探索回数  :', i
    print u'取得したルート:', minimum_route[::-1]
    print u'ルートの距離 :', distance_dict[goal_vtx]
    print u'総実行時間  :', elapsed_time
    
dijkstras()