hashcode-2018/main.go

package main

import (
	"fmt"
	"math/rand"
	"os"
	"sort"
)

var input *os.File
var output *os.File

var R int
var C int
var F int
var N int
var B int
var T int

var Rides []*Ride

var Cars []*Car

var Sched Scheduler

type Ride struct {
	ID int

	a, b, x, y, s, f int

	used bool
}

func (r Ride) Length() int {
	xdist := r.a - r.x
	if xdist < 0 {
		xdist = -xdist
	}
	ydist := r.b - r.y
	if ydist < 0 {
		ydist = -ydist
	}

	return xdist + ydist
}

func abs(x int) int {
	if x < 0 {
		return -x
	} else {
		return x
	}
}

func (r *Ride) length() int {
	return abs(r.a-r.x) + abs(r.b-r.y)
}

type ByEndtime []*Ride

func (rs ByEndtime) Len() int      { return len(rs) }
func (rs ByEndtime) Swap(i, j int) { rs[i], rs[j] = rs[j], rs[i] }
func (rs ByEndtime) Less(i, j int) bool {
	return rs[i].f < rs[j].f || (rs[i].f == rs[j].f && rs[i].length() < rs[j].length())
}

type Scheduler interface {
	Add(*Car)
	RemoveAtIndex(k int) *Car
	Pop() *Car
}

type Car struct {
	ID    int
	Rides []int

	Arrival int
	X       int
	Y       int

	score   int
	pqindex int // used for removal in prority queue
}

func (c *Car) Update(r *Ride) {
	c.moveTo(r.a, r.b)
	if c.Arrival <= r.s {
		c.Arrival = r.s
		// count bonus points for being on time
		c.score += B
	}
	c.moveTo(r.x, r.y)
	c.score += r.length()
	c.Rides = append(c.Rides, r.ID)
}

func (c *Car) EarliestFinish(r *Ride) int {
	copy := &Car{
		Arrival: c.Arrival,
		X:       c.X,
		Y:       c.Y,
	}
	copy.moveTo(r.a, r.b)
	if copy.Arrival < r.s {
		copy.Arrival = r.s
	}
	copy.moveTo(r.x, r.y)
	return copy.Arrival
}

func (c *Car) moveTo(x, y int) {
	xdist := c.X - x
	if xdist < 0 {
		xdist = -xdist
	}
	ydist := c.Y - y
	if ydist < 0 {
		ydist = -ydist
	}

	c.Arrival += xdist + ydist
	c.X = x
	c.Y = y
}

func (c *Car) distanceTo(x, y int) int {
	return abs(c.X-x) + abs(c.Y-y)
}

func max(a, b int) int {
	if a > b {
		return a
	} else {
		return b
	}
}

var bestTotalScore int

func save() {
	if bestTotalScore < 10900000 {
		// TODO find a better way to avoid too many writes
		return
	}
	output.Truncate(0)
	output.Seek(0, 0)
	for _, c := range Cars {
		fmt.Fprintf(output, "%d", len(c.Rides))
		for _, ri := range c.Rides {
			fmt.Fprintf(output, " %d", ri)
		}
		fmt.Fprintf(output, "\n")
	}
	fmt.Printf("%d\n", bestTotalScore)
}

func (c *Car) AssignRideRecur(r *Ride, cumulativeScore int, depth int, fromDepth int) (int, bool) {
	r.used = true
	oldC := *c
	c.Update(r)
	Sched.Add(c)
	// recursion 101
	rdepth, fix := Choose(cumulativeScore+c.score-oldC.score, depth+1, fromDepth)
	Sched.RemoveAtIndex(c.pqindex)
	// pqindex is meaningless now but it's ok
	// it will be fixed when c is added again
	*c = oldC
	r.used = false
	return rdepth, fix
}

func (c *Car) CanPick(r *Ride) bool {
	return !r.used && r.Length() <= 8000 && r.f >= c.EarliestFinish(r)
}

func (c *Car) pickBestRide() *Ride {
	var bestRide *Ride
	var bestRideTotal int
	var bestRideLen int

	for _, r := range Rides {
		if !c.CanPick(r) {
			continue
		}

		lenOfRide := r.length()
		total := max(c.distanceTo(r.a, r.b), r.s-c.Arrival) + lenOfRide
		if bestRide == nil || lenOfRide*bestRideTotal > total*bestRideLen {
			bestRide = r
			bestRideLen = lenOfRide
			bestRideTotal = total
		}
	}

	return bestRide
}

func (c *Car) pickRandomRide() *Ride {
	count := 0
	for _, r := range Rides {
		if !c.CanPick(r) {
			continue
		}
		count++
	}
	i := rand.Intn(count)
	for _, r := range Rides {
		if !c.CanPick(r) {
			continue
		}
		if i == 0 {
			return r
		}
		i = i - 1
	}
	return nil
}

// invariant : when entering and leaving function,
// Sched has the "same" heap. Same means successive
// popped values will be the same (but the internal
// ordering of nodes may be different)
// Return value (a,b) :
// - rdepth >= 0 return back to that depth
// - rdepth < 0 completely unwind stack and finish
// - fix == false try to make a change at depth `rdepth'
// - fix == true change back, use "optimal" ride again
func Choose(cumulativeScore int, depth int, fromDepth int) (int, bool) {
	c := Sched.Pop()
	if c == nil {
		// At this point we have a complete configuration
		fmt.Printf("score obtained: %d depth: %d\n", cumulativeScore, depth)
		if cumulativeScore >= bestTotalScore {
			bestTotalScore = cumulativeScore
			// Go back to random depth and make a new change
			return rand.Intn(depth), false
		} else {
			// Go back to fix change that lowered our score
			return fromDepth, true
		}
	}

	var rdepth int
	var fix bool
	r := c.pickBestRide()
	if r != nil {
		rdepth, fix = c.AssignRideRecur(r, cumulativeScore, depth, fromDepth)
		// if depth == reverse depth, make a change
		// and go back up the stack
		for depth == rdepth {
			var r *Ride
			if fix {
				r = c.pickBestRide()
			} else {
				r = c.pickRandomRide()
			}
			// note fromDepth reset to depth
			rdepth, fix = c.AssignRideRecur(r, cumulativeScore, depth, depth)
		}
	} else {
		// another car may still have other rides
		rdepth, fix = Choose(cumulativeScore, depth, fromDepth)
	}
	// add back the one we popped at beginning of function
	Sched.Add(c)
	// go back down the stack
	return rdepth, fix
}

func solve() {
	rand.Seed(1)

	sort.Sort(ByEndtime(Rides))

	Sched = &prioq{}

	// create cars
	for i := 0; i < F; i++ {
		c := &Car{
			ID:      i,
			Arrival: 0,
			X:       0,
			Y:       0,
		}
		Cars = append(Cars, c)
		Sched.Add(c)
	}

	// start recursion
	Choose(0, 0, 0)
}

func main() {
	sample := os.Args[1]
	fileIn := sample + ".in"
	fileOut := sample + ".out"

	var err error
	input, err = os.Open(fileIn)
	if err != nil {
		panic(fmt.Sprintf("open %s: %v", fileIn, err))
	}
	output, err = os.Create(fileOut)
	if err != nil {
		panic(fmt.Sprintf("creating %s: %v", fileOut, err))
	}
	defer input.Close()
	defer output.Close()

	// Global
	R = readInt()
	C = readInt()
	F = readInt()
	N = readInt()
	B = readInt()
	T = readInt()

	for i := 0; i < N; i++ {
		Rides = append(Rides, &Ride{
			ID: i,
			a:  readInt(),
			b:  readInt(),
			x:  readInt(),
			y:  readInt(),
			s:  readInt(),
			f:  readInt(),
		})
	}

	solve()
}

func readInt() int {
	var i int
	fmt.Fscanf(input, "%d", &i)
	return i
}

func readString() string {
	var str string
	fmt.Fscanf(input, "%s", &str)
	return str
}

func readFloat() float64 {
	var x float64
	fmt.Fscanf(input, "%f", &x)
	return x
}

// Prioq

// Invariant: both children are bigger

type prioq struct {
	bintree []*Car
}

func (c *Car) greater(c2 *Car) bool {
	if c.Arrival == c2.Arrival {
		return c.ID > c2.ID
	}
	return c.Arrival > c2.Arrival
}

func (pq *prioq) Add(car *Car) {
	pq.bintree = append(pq.bintree, car)
	pq.bintree[len(pq.bintree)-1].pqindex = len(pq.bintree) - 1

	// Rebalance tree to respect invariant
	var i = len(pq.bintree) - 1
	var p = (i - 1) / 2
	for p >= 0 && pq.bintree[p].greater(pq.bintree[i]) {
		pq.bintree[p], pq.bintree[i] = pq.bintree[i], pq.bintree[p]
		pq.bintree[p].pqindex = p
		pq.bintree[i].pqindex = i
		i = p
		p = (i - 1) / 2
	}
}

func (pq *prioq) Pop() *Car {
	if len(pq.bintree) == 0 {
		return nil
	}

	if len(pq.bintree) == 1 {
		elem := pq.bintree[0]
		pq.bintree = pq.bintree[:0]
		return elem
	}

	pq.bintree[0].pqindex = -1
	elem := pq.bintree[0]
	// Put last element at root
	pq.bintree[0] = pq.bintree[len(pq.bintree)-1]
	pq.bintree[0].pqindex = 0
	// Remove last element
	pq.bintree = pq.bintree[:len(pq.bintree)-1]

	//        1                  9
	//    10     9	         10     12
	//  11 12   13 14  ->  11 12   13 14
	// 12

	// Rebalance tree to respect invariant
	len := len(pq.bintree)
	i, left, right := 0, 0, 0
	for {
		left = 2*i + 1
		right = 2*i + 2
		if left < len && right < len { // Two children
			if !pq.bintree[left].greater(pq.bintree[right]) {
				if !pq.bintree[i].greater(pq.bintree[left]) {
					break // Inferior to both children
				} else {
					pq.bintree[i], pq.bintree[left] = pq.bintree[left], pq.bintree[i]
					pq.bintree[i].pqindex = i
					pq.bintree[left].pqindex = left
					i = left
				}
			} else {
				if !pq.bintree[i].greater(pq.bintree[right]) {
					break // Inferior to both children
				} else {
					pq.bintree[i], pq.bintree[right] = pq.bintree[right], pq.bintree[i]
					pq.bintree[i].pqindex = i
					pq.bintree[right].pqindex = right
					i = right
				}
			}
		} else if left < len { // One child (left)
			if !pq.bintree[i].greater(pq.bintree[left]) {
				break // Inferior to only child
			}
			pq.bintree[i], pq.bintree[left] = pq.bintree[left], pq.bintree[i]
			pq.bintree[i].pqindex = i
			pq.bintree[left].pqindex = left
			i = left
		} else { // No child
			break
		}

	}

	return elem
}

func (pq *prioq) RemoveAtIndex(k int) *Car {
	if k == 0 {
		return pq.Pop()
	}
	pq.bintree[k].pqindex = -1
	elem := pq.bintree[k]
	reassign := pq.bintree[k+1:]
	pq.bintree = pq.bintree[:k]
	for _, c := range reassign {
		pq.Add(c)
	}
	return elem
}

func (pq *prioq) empty() bool {
	return len(pq.bintree) == 0
}