Jane Street problem

main.go

// Solves this problem: http://1mage.us/1339.

package main

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

	"github.com/unixpickle/approb"
	"github.com/unixpickle/num-analysis/linalg"
)

func main() {
	ls := &LineSet{}
	addTriangle(ls, 7, true, true, true, 0, 0, 1)

	fmt.Println(approb.Covariances(20000000, func() linalg.Vector {
		x, y := ls.Sample()
		return []float64{x, y}
	}))

	// Produces a picture representing the distribution.
	//
	//    img := image.NewGray(image.Rect(0, 0, 1000, 1000))
	//    for i := 0; i < 10000; i++ {
	// 	   x, y := ls.Sample()
	// 	   img.SetGray(int(x*1000), 1000+int(y*1000), color.Gray{0xff})
	//    }
	//    f, _ := os.Create("output.png")
	//    png.Encode(f, img)
	//
}

func addTriangle(ls *LineSet, maxDepth int, doLeft, doRight, doBottom bool, x, y, size float64) {
	height := size * math.Sqrt(3) / 2
	if doLeft {
		ls.Add(&Line{X1: x, X2: x + size/2, Y1: y, Y2: y - height})
	}
	if doRight {
		ls.Add(&Line{X1: x + size, X2: x + size/2, Y1: y, Y2: y - height})
	}
	if doBottom {
		ls.Add(&Line{X1: x, X2: x + size, Y1: y, Y2: y})
	}
	if maxDepth > 0 {
		addTriangle(ls, maxDepth-1, false, true, false, x, y, size/2)
		addTriangle(ls, maxDepth-1, true, false, false, x+size/2, y, size/2)
		addTriangle(ls, maxDepth-1, false, false, true, x+size/4, y-height/2, size/2)
	}
}

type Line struct {
	X1, X2 float64
	Y1, Y2 float64
}

func (l *Line) Length() float64 {
	dx := l.X2 - l.X1
	dy := l.Y2 - l.Y1
	return math.Sqrt(dx*dx + dy*dy)
}

func (l *Line) Sample() (x, y float64) {
	dist := rand.Float64()
	return dist*l.X1 + (1-dist)*l.X2, dist*l.Y1 + (1-dist)*l.Y2
}

type LineSet struct {
	Lines       []*Line
	Offsets     []float64
	TotalLength float64
}

func (l *LineSet) Add(line *Line) {
	l.Lines = append(l.Lines, line)
	l.TotalLength += line.Length()
	l.Offsets = append(l.Offsets, l.TotalLength)
}

func (l *LineSet) Sample() (x, y float64) {
	offset := rand.Float64() * l.TotalLength
	idx := sort.Search(len(l.Offsets), func(idx int) bool {
		return l.Offsets[idx] >= offset || idx == len(l.Offsets)
	})
	return l.Lines[idx].Sample()
}