leetcode146周赛-5132-颜色最短的交替路径

class Solution:
    def shortestAlternatingPaths(self, n: int, red_edges, blue_edges):
        def function(n,r,b):
            result=[[float("inf")]*2 for _ in range(n)]
            result[0]=[0,0]
            r.sort()
            b.sort()
            rdict={}
            bdict={}
            state=0
            for i in range(len(r)):
                if r[i][0] not in rdict:
                    rdict[r[i][0]]=[r[i][1]]
                else:
                    rdict[r[i][0]].append(r[i][1])
            for i in range(len(b)):
                if b[i][0] not in bdict:
                    bdict[b[i][0]]=[b[i][1]]
                else:
                    bdict[b[i][0]].append(b[i][1])
            while state==0:
                state=1
                for i in range(n):
                    if i in rdict:
                        for j in rdict[i]:
                            if result[j][0]>result[i][1]+1:
                                state=0
                            result[j][0]=min(result[j][0],result[i][1]+1)
                    if i in bdict:
                        for j in bdict[i]:
                            if result[j][1]>result[i][0]+1:
                                state=0
                            result[j][1]=min(result[j][1],result[i][0]+1)
            tmpresult=[]
            for i in range(len(result)):
                if min(result[i])==float("inf"):
                    tmpresult.append(-1)
                else:
                    tmpresult.append(min(result[i]))
            return tmpresult
        return function(n,red_edges,blue_edges)