单源最短路径

  1. import scala.reflect.ClassTag
  3. import org.apache.spark.graphx._
  5. /**
  6.  * Computes shortest paths to the given set of landmark vertices, returning a graph where each
  7.  * vertex attribute is a map containing the shortest-path distance to each reachable landmark.
  8.  */
  9. object ShortestPaths {
  10.   /** Stores a map from the vertex id of a landmark to the distance to that landmark. */
  11.   type SPMap = Map[VertexId, Int]
  13.   private def makeMap(x: (VertexId, Int)*) = Map(x: _*)
  15.   private def incrementMap(spmap: SPMap): SPMap = spmap.map { case (v, d) => v -> (d + 1) }
  17.   private def addMaps(spmap1: SPMap, spmap2: SPMap): SPMap =
  18.     (spmap1.keySet ++ spmap2.keySet).map {
  19.       k => k -> math.min(spmap1.getOrElse(k, Int.MaxValue), spmap2.getOrElse(k, Int.MaxValue))
  20.     }.toMap
  22.   /**
  23.    * Computes shortest paths to the given set of landmark vertices.
  24.    *
  25.    * @tparam ED the edge attribute type (not used in the computation)
  26.    *
  27.    * @param graph the graph for which to compute the shortest paths
  28.    * @param landmarks the list of landmark vertex ids. Shortest paths will be computed to each
  29.    * landmark.
  30.    *
  31.    * @return a graph where each vertex attribute is a map containing the shortest-path distance to
  32.    * each reachable landmark vertex.
  33.    */
  34.   def run[VD, ED: ClassTag](graph: Graph[VD, ED], landmarks: Seq[VertexId]): Graph[SPMap, ED] = {
  35.     val spGraph = graph.mapVertices { (vid, attr) =>
  36.       if (landmarks.contains(vid)) makeMap(vid -> 0) else makeMap()
  37.     }
  39.     val initialMessage = makeMap()
  41.     def vertexProgram(id: VertexId, attr: SPMap, msg: SPMap): SPMap = {
  42.       addMaps(attr, msg)
  43.     }
  45.     def sendMessage(edge: EdgeTriplet[SPMap, _]): Iterator[(VertexId, SPMap)] = {
  46.       val newAttr = incrementMap(edge.dstAttr)
  47.       if (edge.srcAttr != addMaps(newAttr, edge.srcAttr)) Iterator((edge.srcId, newAttr))
  48.       else Iterator.empty
  49.     }
  51.     Pregel(spGraph, initialMessage)(vertexProgram, sendMessage, addMaps)
  52.   }
  53. }