"Optimized" Water Disruption solution

water2.go

package main

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"log"
	"math/rand"
	"os"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
)

var debug = false

func DoWater(input io.Reader) {
	scanner := bufio.NewScanner(input) // was os.Stdin
	scanner.Scan()
	nlines, _ := strconv.Atoi(scanner.Text())

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

	for i := 0; i < nlines; i++ {
		scanner.Scan()
		line := scanner.Text()
		words := strings.Split(line, " ")
		costs[i] = make([]int16, nlines)
		for k, word := range words {
			bigint, _ := strconv.Atoi(word)
			costs[i][k] = int16(bigint)
		}
	}

	scanner.Scan()
	line := scanner.Text()
	words := strings.Split(line, " ")
	newWellCosts := make([]int16, nlines)
	for k, word := range words {
		bigint, _ := strconv.Atoi(word)
		newWellCosts[k] = int16(bigint)
	}

	areasLeft := make([]int, nlines)
	for i, _ := range areasLeft {
		areasLeft[i] = i
	}

	// Similar problem as the friends stuff.
	// Note that this is pretty much unused except for info during debugging:
	connectedAreas := make([][]int16, 0)
	connectedAreasLen := int16(0)
	connectedAreasMap := make([]int16, len(areasLeft))
	for i, _ := range connectedAreasMap {
		connectedAreasMap[i] = -1
	}

	totalCosts := 0
	numAreasLeft := len(areasLeft)
	for numAreasLeft > 0 {
		connectionCost := int16(32767)
		toArea := -1
		fromAreaLeftIdx := -1
		fromArea := -1

		// Find lowest cost action for connecting an area to an established well
		for i, area := range areasLeft {
			if area < 0 {
				continue
			}
			for k, camk := range connectedAreasMap {
				if camk >= 0 {
					if costs[area][k] < connectionCost {
						connectionCost = costs[area][k]
						fromArea = k
						toArea = area
						fromAreaLeftIdx = i
					}
				}
			}
		}

		// Find building new well costs
		buildNewArea := -1
		buildNewCost := int16(32767)
		buildNewAreaLeftIdx := -1
		for k, area := range areasLeft {
			if area < 0 {
				continue
			}
			if newWellCosts[area] < buildNewCost {
				buildNewArea = area
				buildNewCost = newWellCosts[area]
				buildNewAreaLeftIdx = k
			}
		}

		if toArea < 0 || buildNewCost < connectionCost {
			if debug {
				fmt.Println("Build new well at", buildNewArea, "at cost", buildNewCost)
				connectedAreas = append(connectedAreas, []int16{int16(buildNewArea)})
				connectedAreasMap[buildNewArea] = int16(len(connectedAreas) - 1)
			} else {
				connectedAreasLen++
				connectedAreasMap[buildNewArea] = connectedAreasLen
			}
			areasLeft[buildNewAreaLeftIdx] = -1
			numAreasLeft--
			totalCosts += int(buildNewCost)
		} else {
			if debug {
				fmt.Println("Connect water from", fromArea, "to", toArea, "at cost", connectionCost)
				connectedAreas[connectedAreasMap[fromArea]] = append(connectedAreas[connectedAreasMap[fromArea]], int16(toArea))
			}

			connectedAreasMap[toArea] = connectedAreasMap[fromArea]
			areasLeft[fromAreaLeftIdx] = -1
			numAreasLeft--
			totalCosts += int(connectionCost)
		}
		if debug {
			fmt.Println(areasLeft)
			fmt.Println(connectedAreas) // Graph colors. (??? terminology)
		}
	}
	fmt.Println(totalCosts)
}

func generateWaterPuzzle(numNodes int, seed int64) io.Reader {
	nodes := make([][]int, numNodes)
	for n, _ := range nodes {
		nodes[n] = make([]int, numNodes)
	}
	r := rand.New(rand.NewSource(seed))
	for n, _ := range nodes {
		for k, _ := range nodes[n] {
			rnd := r.Intn(1000) + 1
			nodes[k][n] = rnd
			nodes[n][k] = nodes[k][n]
		}
	}
	newWell := make([]int, numNodes)
	for n, _ := range nodes {
		rnd := r.Intn(1000) + 1
		newWell[n] = rnd
	}
	s := strconv.Itoa(numNodes) + "\n"

	for _, n2 := range nodes {
		k := fmt.Sprintln(n2)
		s += k[1:len(k)-2] + "\n"
	}

	k := fmt.Sprintln(newWell)
	s += k[1:len(k)-2] + "\n"
	if debug {
		fmt.Println(s)
	}
	return bytes.NewBufferString(s)
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	// Log CPU usage profile
	cpuprofile := "cpu.pprof"
	f, err := os.Create(cpuprofile)
	if err != nil {
		log.Fatal(err)
	}
	pprof.StartCPUProfile(f)
	defer pprof.StopCPUProfile()

	// Log memory usage profile
	memprofile := "mem.pprof"
	f2, err := os.Create(memprofile)
	if err != nil {
		log.Fatal(err)
	}
	defer pprof.WriteHeapProfile(f2)

	DoWater(bytes.NewBufferString(`4
0 2 2 2
2 0 3 3
2 3 0 4
2 3 4 0
5 4 4 3`))
	for i := 1; i <= 20; i++ {
		DoWater(generateWaterPuzzle(500, int64(i)))
	}
	// Currently: 5.8 seconds for 1 example solve, plus 20 random maximum width cases
	// 2160.03kB mem usage go 1.4.2
}