github.com/livekit/protocol@v1.16.1-0.20240517185851-47e4c6bba773/utils/graph.go (about) 1 // Copyright 2023 LiveKit, Inc. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package utils 16 17 import ( 18 "container/heap" 19 "log" 20 "math" 21 22 "github.com/gammazero/deque" 23 ) 24 25 type GraphNodeProps[K comparable] interface { 26 ID() K 27 } 28 29 type GraphEdgeProps interface { 30 Length() int64 31 } 32 33 type SimpleGraphEdge struct{} 34 35 func (e SimpleGraphEdge) Length() int64 { return 1 } 36 37 type Graph[K comparable, N GraphNodeProps[K], E GraphEdgeProps] struct { 38 nodesByID map[K]*GraphNode[N] 39 freeIndices *deque.Deque[int] 40 nodes []*GraphNode[N] 41 edges [][]*GraphEdge[N, E] 42 } 43 44 func NewGraph[K comparable, N GraphNodeProps[K], E GraphEdgeProps]() *Graph[K, N, E] { 45 return &Graph[K, N, E]{ 46 nodesByID: map[K]*GraphNode[N]{}, 47 freeIndices: deque.New[int](0), 48 } 49 } 50 51 func (g *Graph[K, N, E]) Size() int { 52 return len(g.nodes) 53 } 54 55 func (g *Graph[K, N, E]) NodeIDs() []K { 56 ids := make([]K, 0, len(g.nodes)-g.freeIndices.Len()) 57 for _, n := range g.nodes { 58 if n != nil { 59 ids = append(ids, n.props.ID()) 60 } 61 } 62 return ids 63 } 64 65 func (g *Graph[K, N, E]) InsertNode(props N) { 66 if n, ok := g.nodesByID[props.ID()]; ok { 67 n.props = props 68 return 69 } 70 71 var i int 72 if g.freeIndices.Len() != 0 { 73 i = g.freeIndices.PopBack() 74 } else { 75 i = len(g.nodes) 76 g.nodes = append(g.nodes, nil) 77 for j := range g.edges { 78 g.edges[j] = append(g.edges[j], nil) 79 } 80 g.edges = append(g.edges, make([]*GraphEdge[N, E], len(g.nodes))) 81 } 82 83 n := &GraphNode[N]{ 84 i: i, 85 props: props, 86 } 87 88 g.nodes[i] = n 89 g.nodesByID[props.ID()] = n 90 } 91 92 func (g *Graph[K, N, E]) DeleteNode(id K) { 93 n, ok := g.nodesByID[id] 94 if !ok { 95 return 96 } 97 98 delete(g.nodesByID, id) 99 g.nodes[n.i] = nil 100 101 for _, es := range g.edges { 102 es[n.i] = nil 103 } 104 for j := range g.edges[n.i] { 105 g.edges[n.i][j] = nil 106 } 107 108 g.freeIndices.PushBack(n.i) 109 } 110 111 func (g *Graph[K, N, E]) InsertEdge(src, dst K, props E) { 112 s := g.nodesByID[src] 113 d := g.nodesByID[dst] 114 115 g.edges[s.i][d.i] = &GraphEdge[N, E]{props} 116 } 117 118 func (g *Graph[K, N, E]) DeleteEdge(src, dst K) { 119 s := g.nodesByID[src] 120 d := g.nodesByID[dst] 121 122 g.edges[s.i][d.i] = nil 123 } 124 125 func (g *Graph[K, N, E]) HasNode(id K) bool { 126 return g.nodesByID[id] != nil 127 } 128 129 func (g *Graph[K, N, E]) Node(id K) (props N) { 130 n := g.nodesByID[id] 131 if n == nil { 132 return 133 } 134 return n.props 135 } 136 137 func (g *Graph[K, N, E]) HasEdge(src, dst K) bool { 138 s := g.nodesByID[src] 139 d := g.nodesByID[dst] 140 if s == nil || d == nil { 141 return false 142 } 143 144 return g.edges[s.i][d.i] != nil 145 } 146 147 func (g *Graph[K, N, E]) Edge(src, dst K) (p E) { 148 s := g.nodesByID[src] 149 d := g.nodesByID[dst] 150 if s == nil || d == nil { 151 return 152 } 153 154 e := g.edges[s.i][d.i] 155 if e == nil { 156 return 157 } 158 return e.props 159 } 160 161 func (g *Graph[K, N, E]) OutEdges(src K) map[K]E { 162 s := g.nodesByID[src] 163 if s == nil { 164 return nil 165 } 166 167 edges := make(map[K]E, len(g.nodes)) 168 for i, e := range g.edges[s.i] { 169 if e != nil { 170 edges[g.nodes[i].props.ID()] = e.props 171 } 172 } 173 return edges 174 } 175 176 func (g *Graph[K, N, E]) InEdges(dst K) map[K]E { 177 d := g.nodesByID[dst] 178 if d == nil { 179 return nil 180 } 181 182 edges := make(map[K]E, len(g.nodes)) 183 for i, es := range g.edges { 184 if es[d.i] != nil { 185 edges[g.nodes[i].props.ID()] = es[d.i].props 186 } 187 } 188 return edges 189 } 190 191 func (g *Graph[K, N, E]) ShortestPath(src, dst K) ([]N, int64) { 192 paths := &graphPathMinHeap[N]{} 193 visited := map[*GraphNode[N]]*graphPath[N]{} 194 195 s := g.nodesByID[src] 196 d := g.nodesByID[dst] 197 if s == nil || d == nil { 198 return nil, 0 199 } 200 201 path := &graphPath[N]{node: s} 202 heap.Push(paths, path) 203 visited[path.node] = path 204 205 for { 206 if paths.Len() == 0 { 207 return nil, 0 208 } 209 210 prev := heap.Pop(paths).(*graphPath[N]) 211 for i, e := range g.edges[prev.node.i] { 212 if e == nil { 213 continue 214 } 215 216 path := &graphPath[N]{ 217 prev: prev, 218 node: g.nodes[i], 219 length: prev.length + e.props.Length(), 220 num: prev.num + 1, 221 } 222 223 if p, ok := visited[path.node]; ok && p.Less(path) { 224 continue 225 } 226 visited[path.node] = path 227 228 if path.node == d { 229 return path.Nodes(), path.length 230 } 231 232 heap.Push(paths, path) 233 } 234 } 235 } 236 237 func (g *Graph[K, N, E]) TopologicalSort() []N { 238 if g.Size() == 0 { 239 return nil 240 } 241 242 log.Println(len(g.nodes)) 243 nodes := make([]N, 0, len(g.nodes)) 244 acyclic := true 245 246 temporary := make(map[*GraphNode[N]]struct{}, len(g.nodes)) 247 permanent := make(map[*GraphNode[N]]struct{}, len(g.nodes)) 248 249 for _, n := range g.nodes { 250 if _, ok := permanent[n]; ok { 251 continue 252 } 253 254 g.traverseDepthFirst(n, func(n *GraphNode[N], next func()) { 255 if _, ok := permanent[n]; ok { 256 return 257 } 258 if _, ok := temporary[n]; ok { 259 acyclic = false 260 return 261 } 262 263 temporary[n] = struct{}{} 264 265 next() 266 267 delete(temporary, n) 268 permanent[n] = struct{}{} 269 nodes = append(nodes, n.props) 270 }) 271 } 272 273 if !acyclic { 274 return nil 275 } 276 277 for i := 0; i < len(nodes)/2; i++ { 278 nodes[i], nodes[len(nodes)-1-i] = nodes[len(nodes)-1-i], nodes[i] 279 } 280 return nodes 281 } 282 283 func (g *Graph[K, N, E]) traverseDepthFirst(n *GraphNode[N], fn func(n *GraphNode[N], next func())) { 284 fn(n, func() { 285 for i, e := range g.edges[n.i] { 286 if e != nil { 287 g.traverseDepthFirst(g.nodes[i], fn) 288 } 289 } 290 }) 291 } 292 293 type graphPath[T any] struct { 294 prev *graphPath[T] 295 node *GraphNode[T] 296 length int64 297 num int 298 } 299 300 func (p *graphPath[T]) nodes(i int) []T { 301 if p.prev == nil { 302 return append(make([]T, 0, i), p.node.props) 303 } else { 304 return append(p.prev.nodes(i+1), p.node.props) 305 } 306 } 307 308 func (p *graphPath[T]) Nodes() []T { 309 return p.nodes(1) 310 } 311 312 func (p *graphPath[T]) Less(o *graphPath[T]) bool { 313 return (p.length == o.length && p.num < o.num) || p.length < o.length 314 } 315 316 type graphPathMinHeap[T any] []*graphPath[T] 317 318 func (h *graphPathMinHeap[T]) Len() int { 319 return len(*h) 320 } 321 322 func (h *graphPathMinHeap[T]) Less(i, j int) bool { 323 return (*h)[i].Less((*h)[j]) 324 } 325 326 func (h *graphPathMinHeap[T]) Swap(i, j int) { 327 (*h)[i], (*h)[j] = (*h)[j], (*h)[i] 328 } 329 330 func (h *graphPathMinHeap[T]) Push(x any) { 331 *h = append(*h, x.(*graphPath[T])) 332 } 333 334 func (h *graphPathMinHeap[T]) Pop() any { 335 x := (*h)[len(*h)-1] 336 (*h)[len(*h)-1] = nil 337 *h = (*h)[:len(*h)-1] 338 return x 339 } 340 341 type GraphNode[T any] struct { 342 i int 343 props T 344 } 345 346 type GraphEdge[N, E any] struct { 347 props E 348 } 349 350 const inf = int64(math.MaxInt64/2 - 1) 351 352 func NewFlowGraph(n int64) FlowGraph { 353 cap := make([]int64, n*n) 354 cost := make([]int64, n*n) 355 return FlowGraph{n, cap, cost} 356 } 357 358 type FlowGraph struct { 359 n int64 360 cap, cost []int64 361 } 362 363 func (g *FlowGraph) AddEdge(s, t, cap, cost int64) { 364 g.cap[s*g.n+t] = cap 365 g.cap[t*g.n+s] = cap 366 g.cost[s*g.n+t] = cost 367 g.cost[t*g.n+s] = cost 368 } 369 370 type MinCostMaxFlow struct { 371 found []bool 372 n int64 373 cap, flow, cost []int64 374 prev, dist, pi []int64 375 } 376 377 func (f *MinCostMaxFlow) search(s, t int64) bool { 378 for i := range f.found { 379 f.found[i] = false 380 } 381 for i := range f.dist { 382 f.dist[i] = inf 383 } 384 385 f.dist[s] = 0 386 387 for s != f.n { 388 best := f.n 389 f.found[s] = true 390 391 for i := int64(0); i < f.n; i++ { 392 if f.found[i] { 393 continue 394 } 395 396 if f.flow[i*f.n+s] != 0 { 397 val := f.dist[s] + f.pi[s] - f.pi[i] - f.cost[i*f.n+s] 398 if f.dist[i] > val { 399 f.dist[i] = val 400 f.prev[i] = s 401 } 402 } 403 404 if f.flow[s*f.n+i] < f.cap[s*f.n+i] { 405 val := f.dist[s] + f.pi[s] - f.pi[i] + f.cost[s*f.n+i] 406 if f.dist[i] > val { 407 f.dist[i] = val 408 f.prev[i] = s 409 } 410 } 411 412 if f.dist[i] < f.dist[best] { 413 best = i 414 } 415 } 416 417 s = best 418 } 419 420 for i := int64(0); i < f.n; i++ { 421 pi := f.pi[i] + f.dist[i] 422 if pi > inf { 423 pi = inf 424 } 425 f.pi[i] = pi 426 } 427 428 return f.found[t] 429 } 430 431 func (f *MinCostMaxFlow) Flow(s, t int64) int64 { 432 return f.flow[s*f.n+t] 433 } 434 435 func (f *MinCostMaxFlow) ComputeMaxFlow(g FlowGraph, s, t int64) (flow, cost int64) { 436 f.cap = g.cap 437 f.cost = g.cost 438 f.n = g.n 439 440 f.found = make([]bool, f.n) 441 f.flow = make([]int64, f.n*f.n) 442 f.dist = make([]int64, f.n+1) 443 f.prev = make([]int64, f.n) 444 f.pi = make([]int64, f.n) 445 446 for f.search(s, t) { 447 pathFlow := inf 448 for u := t; u != s; u = f.prev[u] { 449 var pf int64 450 if f.flow[u*f.n+f.prev[u]] != 0 { 451 pf = f.flow[u*f.n+f.prev[u]] 452 } else { 453 pf = f.cap[f.prev[u]*f.n+u] - f.flow[f.prev[u]*f.n+u] 454 } 455 if pf < pathFlow { 456 pathFlow = pf 457 } 458 } 459 460 for u := t; u != s; u = f.prev[u] { 461 if f.flow[u*f.n+f.prev[u]] != 0 { 462 f.flow[u*f.n+f.prev[u]] -= pathFlow 463 cost -= pathFlow * f.cost[u*f.n+f.prev[u]] 464 } else { 465 f.flow[f.prev[u]*f.n+u] += pathFlow 466 cost += pathFlow * f.cost[f.prev[u]*f.n+u] 467 } 468 } 469 flow += pathFlow 470 } 471 472 return flow, cost 473 }