Code kata. https://forum.lowyat.net/index.php?showtopic=3571866&view=findpost&p=75952192

water.go

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"regexp"
	"strconv"
)

var debug = true

func doHuffman(input io.Reader) {
	re, _ := regexp.Compile("(?i)[0-9]+")

	scanner := bufio.NewScanner(input) // was os.Stdin
	scanner.Scan()
	// I'm getting lazy
	nlines, _ := strconv.Atoi(scanner.Text())

	costs := make([][]int, nlines)

	for i := 0; i < nlines; i++ {
		scanner.Scan()
		line := scanner.Text()
		words := re.FindAllStringSubmatch(line, -1)
		costs[i] = make([]int, nlines)
		for k, word := range words {
			costs[i][k], _ = strconv.Atoi(word[0])
		}
	}

	scanner.Scan()
	line := scanner.Text()
	words := re.FindAllStringSubmatch(line, -1)
	newWellCosts := make([]int, nlines)
	for k, word := range words {
		newWellCosts[k], _ = strconv.Atoi(word[0])
	}

	//fmt.Println(costs)
	//fmt.Println(newWellCosts)
	areasLeft := make([]int, nlines)
	for i, _ := range areasLeft {
		areasLeft[i] = i
	}

	// Similar problem as the friends stuff.
	connectedAreas := make([][]int, 0)
	connectedAreasMap := make(map[int]int, 0)

	totalCosts := 0
	for len(areasLeft) > 0 {
		c := 99999999
		toArea := -1
		fromAreaLeftIdx := -1
		fromArea := -1
		// Find lowest cost action.
		for i, area := range areasLeft {
			for k, _ := range costs {
				_, connected := connectedAreasMap[k]

				if connected && costs[area][k] < c {
					c = costs[area][k]
					fromArea = k
					toArea = area
					fromAreaLeftIdx = i
				}
			}
		}

		// Find building new well costs
		buildNewArea := -1
		buildNewCost := 999999
		buildNewAreaLeftIdx := -1
		for k, area := range areasLeft {
			if newWellCosts[area] < buildNewCost {
				buildNewArea = area
				buildNewCost = newWellCosts[area]
				buildNewAreaLeftIdx = k
			}
		}

		if toArea < 0 || buildNewCost < c {
			if debug {
				fmt.Println("Build new well at", buildNewArea, "at cost", buildNewCost)
			}

			connectedAreas = append(connectedAreas, []int{buildNewArea})
			connectedAreasMap[buildNewArea] = len(connectedAreas) - 1
			areasLeft = append(areasLeft[:buildNewAreaLeftIdx], areasLeft[buildNewAreaLeftIdx+1:]...)
			totalCosts += buildNewCost
		} else {
			if debug {
				fmt.Println("Connect water from", fromArea, "to", toArea, "at cost", c)
			}
			connectedAreas[connectedAreasMap[fromArea]] = append(connectedAreas[connectedAreasMap[fromArea]], toArea)
			connectedAreasMap[toArea] = connectedAreasMap[fromArea]
			areasLeft = append(areasLeft[:fromAreaLeftIdx], areasLeft[fromAreaLeftIdx+1:]...)
			totalCosts += c
		}
		if debug {
			fmt.Println(areasLeft)
			fmt.Println(connectedAreas) // Graph colors. (??? terminology)
		}
	}
	fmt.Println(totalCosts)
}

func main() {
	input := bytes.NewBufferString(`4
0 2 2 2
2 0 3 3
2 3 0 4
2 3 4 0
5 4 4 3`)
	doHuffman(input)
	if debug {
		fmt.Println("------")
		input = bytes.NewBufferString(`5
0 2 2 2 99
2 0 3 3 99
2 3 0 4 99
2 3 4 0 99
99 99 99 99 99
5 4 4 3 1`)
		doHuffman(input)
	}
}