Permalink
Switch branches/tags
Nothing to show
Find file
Fetching contributors…
Cannot retrieve contributors at this time
56 lines (47 sloc) 1.29 KB
require 'rubygems'
require 'decisiontree'
include DecisionTree
# ---Continuous---
# Read in the training data
training = []
attributes = nil
File.open('data/continuous-training.txt', 'r').each_line do |line|
data = line.strip.chomp('.').split(',')
attributes ||= data
training_data = data.collect do |v|
case v
when 'healthy'
1
when 'colic'
0
else
v.to_f
end
end
training.push(training_data)
end
# Remove the attribute row from the training data
training.shift
# Instantiate the tree, and train it based on the data (set default to '1')
dec_tree = ID3Tree.new(attributes, training, 1, :continuous)
dec_tree.train
# ---Test the tree---
# Read in the test cases
# Note: omit the attribute line (first line), we know the labels from the training data
test = []
File.open('data/continuous-test.txt', 'r').each_line do |line|
data = line.strip.chomp('.').split(',')
test_data = data.collect do |v|
if v == 'healthy' || v == 'colic'
v == 'healthy' ? 1 : 0
else
v.to_f
end
end
test.push(test_data)
end
# Let the tree predict the output and compare it to the true specified value
test.each do |t|
predict = dec_tree.predict(t)
puts "Predict: #{predict} ... True: #{t.last}"
end