decision-tree.go

package main

import (
	"fmt"
	"math"
)

type AttributeValue struct {
	Entropy  float64
	PofValue float64
}

func main() {
	// pos neg attr
	allDataSize := float64(5)
	entropySet := entropy(3, 2, "all")

	// age
	e1 := entropy(1, 0, "<25 age")
	e2 := entropy(1, 0, "25-40 age")
	e3 := entropy(1, 2, ">40 age")
	gain(entropySet,
		AttributeValue{e1, float64(1) / allDataSize},
		AttributeValue{e2, float64(1) / allDataSize},
		AttributeValue{e3, float64(3) / allDataSize},
	)

	// pulse
	e1 = entropy(2, 1, "normal pulse")
	e2 = entropy(1, 1, "rapid pulse")
	gain(entropySet,
		AttributeValue{e1, float64(3) / allDataSize},
		AttributeValue{e2, float64(2) / allDataSize},
	)

	// bp
	e1 = entropy(3, 0, "normal bp")
	e2 = entropy(0, 2, "high bp")
	gain(entropySet,
		AttributeValue{e1, float64(3) / allDataSize},
		AttributeValue{e2, float64(2) / allDataSize},
	)
}

func gain(entropySet float64, attrValues ...AttributeValue) float64 {
	thisGain := entropySet
	formulaString := fmt.Sprintf("%.2f", thisGain)
	for _, av := range attrValues {
		thisGain = thisGain - (av.PofValue * av.Entropy)
		formulaString = fmt.Sprintf("%s - (%.2f * %.2f)", formulaString, av.PofValue, av.Entropy)
	}
	fmt.Println("G =", formulaString)
	fmt.Printf("gain is %.3f\n\n", thisGain)

	return thisGain
}

func entropy(pos, neg float64, attr string) float64 {
	var entropy, PofPos, PofNeg float64
	if pos == 0 || neg == 0 {
		entropy = 0
	} else {
		total := pos + neg
		PofPos = pos / total
		PofNeg = neg / total

		entropy = (-1 * PofPos * math.Log2(PofPos)) - (PofNeg * math.Log2(PofNeg))
	}

	fmt.Printf("E = -1 * (%.2f)(%.2f) - (%.2f)(%.2f)\n", PofPos, math.Log2(PofPos), PofNeg, math.Log2(PofNeg))
	fmt.Printf("entropy for %s: %.2f\n", attr, entropy)
	return entropy
}