Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- package main
- import (
- "container/heap"
- "fmt"
- "math"
- "time"
- "github.com/skorobogatov/input"
- )
- type Vertex struct {
- x, y int
- }
- type Edge struct {
- a, b *setItem
- priority, index int
- }
- type setItem struct {
- parent *setItem
- x *Vertex
- val, depth int
- }
- type PriorityQueue []*Edge
- func (pq PriorityQueue) Len() int { return len(pq) }
- func (pq PriorityQueue) Less(i, j int) bool {
- return pq[i].priority < pq[j].priority
- }
- func (pq PriorityQueue) Swap(i, j int) {
- pq[i], pq[j] = pq[j], pq[i]
- pq[i].index = i
- pq[j].index = j
- }
- func (pq *PriorityQueue) Push(x interface{}) {
- n := len(*pq)
- item := x.(*Edge)
- item.index = n
- *pq = append(*pq, item)
- }
- func (pq *PriorityQueue) Pop() interface{} {
- old := *pq
- n := len(old)
- item := old[n-1]
- item.index = -1
- *pq = old[0 : n-1]
- return item
- }
- func makeSet(x *Vertex) *setItem {
- t := new(setItem)
- t.x = x
- t.parent = t
- t.depth = 0
- return t
- }
- func find(x *setItem) *setItem {
- if x.parent == x {
- return x
- } else {
- x.parent = find(x.parent)
- return x.parent
- }
- }
- func union(x *setItem, y *setItem) {
- rootY := find(y)
- rootX := find(x)
- if rootX.depth < rootY.depth {
- rootX.parent = rootY
- } else {
- rootY.parent = rootX
- if rootX.depth == rootY.depth && rootX != rootY {
- rootX.depth++
- }
- }
- }
- func main() {
- start := time.Now()
- var n, a, b int
- input.Scanf("%d", &n)
- points := make([](*Vertex), n)
- for i := 0; i < n; i++ {
- input.Scanf("%d%d", &a, &b)
- points[i] = &Vertex{a, b}
- }
- pq := make(PriorityQueue, (n*n-n)/2)
- BuildGraph(points, &pq)
- heap.Init(&pq)
- //heap.Init(&pq)
- fmt.Printf("%.2f\n", spanningTree(n, &pq))
- end := time.Now()
- fmt.Println(end.Sub(start))
- }
- func BuildGraph(pointsArr [](*Vertex), pq *PriorityQueue) {
- n := len(pointsArr)
- set := make([](*setItem), n)
- for i, val := range pointsArr {
- set[i] = makeSet(val)
- }
- cnt := 0
- for i := 0; i < n; i++ {
- for j := i + 1; j < n; j++ {
- if i == j {
- continue
- }
- a := ((pointsArr[j].x) - (pointsArr[i].x))
- b := ((pointsArr[j].y) - (pointsArr[i].y))
- dist := ((a * a) + (b * b))
- (*pq)[cnt] = &Edge{set[i], set[j], dist, -1}
- //heap.Push(pq, &Edge{set[i], set[j], dist, -1})
- cnt++
- }
- }
- }
- func spanningTree(n int, edges *PriorityQueue) float64 {
- totalDist := 0.0
- m := len(*edges)
- cnt := 0
- for i := 0; i < m && cnt < (n-1); i++ {
- elem := (heap.Pop(edges).(*Edge))
- if find(elem.a) != find(elem.b) {
- totalDist += math.Sqrt((float64)(elem.priority))
- cnt++
- union(elem.a, elem.b)
- }
- }
- return totalDist
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement