Skip to content
This repository

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
branch: master
Fetching contributors…

Octocat-spinner-32-eaf2f5

Cannot retrieve contributors at this time

file 35 lines (26 sloc) 1.291 kb
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34
require 'rubygems'
require 'decisiontree'

# ---Discrete-----------------------------------------------------------------------------------------

# Read in the training data
training, attributes = [], nil
File.open('data/discrete-training.txt','r').each_line { |line|
  data = line.strip.split(',')
  attributes ||= data
  training.push(data.collect {|v| (v == 'will buy') || (v == "won't buy") ? (v == 'will buy' ? 1 : 0) : v})
}

# 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 = DecisionTree::ID3Tree.new(attributes, training, 1, :discrete)
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/discrete-test.txt','r').each_line { |line| data = line.strip.split(',')
  test.push(data.collect {|v| (v == 'will buy') || (v == "won't buy") ? (v == 'will buy' ? 1 : 0) : v})
}

# Let the tree predict the output and compare it to the true specified value
test.each { |t| predict = dec_tree.predict(t); puts "Predict: #{predict} ... True: #{t.last}"; }

# Graph the tree, save to 'discrete.png'
dec_tree.graph("discrete")
Something went wrong with that request. Please try again.