aima.core.learning.inductive.DecisionTree Maven / Gradle / Ivy

package aima.core.learning.inductive;

import java.util.ArrayList;
import java.util.Hashtable;
import java.util.List;

import aima.core.learning.framework.DataSet;
import aima.core.learning.framework.Example;
import aima.core.util.Util;

/**
 * @author Ravi Mohan
 * 
 */
public class DecisionTree {
	private String attributeName;

	// each node modelled as a hash of attribute_value/decisiontree
	private Hashtable nodes;

	protected DecisionTree() {

	}

	public DecisionTree(String attributeName) {
		this.attributeName = attributeName;
		nodes = new Hashtable();

	}

	public void addLeaf(String attributeValue, String decision) {
		nodes.put(attributeValue, new ConstantDecisonTree(decision));
	}

	public void addNode(String attributeValue, DecisionTree tree) {
		nodes.put(attributeValue, tree);
	}

	public Object predict(Example e) {
		String attrValue = e.getAttributeValueAsString(attributeName);
		if (nodes.containsKey(attrValue)) {
			return nodes.get(attrValue).predict(e);
		} else {
			throw new RuntimeException("no node exists for attribute value "
					+ attrValue);
		}
	}

	public static DecisionTree getStumpFor(DataSet ds, String attributeName,
			String attributeValue, String returnValueIfMatched,
			List unmatchedValues, String returnValueIfUnmatched) {
		DecisionTree dt = new DecisionTree(attributeName);
		dt.addLeaf(attributeValue, returnValueIfMatched);
		for (String unmatchedValue : unmatchedValues) {
			dt.addLeaf(unmatchedValue, returnValueIfUnmatched);
		}
		return dt;
	}

	public static List getStumpsFor(DataSet ds,
			String returnValueIfMatched, String returnValueIfUnmatched) {
		List attributes = ds.getNonTargetAttributes();
		List trees = new ArrayList();
		for (String attribute : attributes) {
			List values = ds.getPossibleAttributeValues(attribute);
			for (String value : values) {
				List unmatchedValues = Util.removeFrom(
						ds.getPossibleAttributeValues(attribute), value);

				DecisionTree tree = getStumpFor(ds, attribute, value,
						returnValueIfMatched, unmatchedValues,
						returnValueIfUnmatched);
				trees.add(tree);

			}
		}
		return trees;
	}

	/**
	 * @return Returns the attributeName.
	 */
	public String getAttributeName() {
		return attributeName;
	}

	@Override
	public String toString() {
		return toString(1, new StringBuffer());
	}

	public String toString(int depth, StringBuffer buf) {

		if (attributeName != null) {
			buf.append(Util.ntimes("\t", depth));
			buf.append(Util.ntimes("***", 1));
			buf.append(attributeName + " \n");
			for (String attributeValue : nodes.keySet()) {
				buf.append(Util.ntimes("\t", depth + 1));
				buf.append("+" + attributeValue);
				buf.append("\n");
				DecisionTree child = nodes.get(attributeValue);
				buf.append(child.toString(depth + 1, new StringBuffer()));
			}
		}

		return buf.toString();
	}
}