com.blazebit.ai.decisiontree.impl.DiscreteDecisionNode Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of blaze-ai-utils Show documentation
Artificial Intelligence Utilities.
There is a newer version: 0.1.21
package com.blazebit.ai.decisiontree.impl;

import com.blazebit.ai.decisiontree.*;

import java.util.*;

/**
 * @author Christian Beikov
 */
public class DiscreteDecisionNode implements DecisionNode {

    private final Attribute attribute;
    private final Map> children;

    public DiscreteDecisionNode(final DecisionNodeFactory decisionNodeFactory, final DiscreteAttribute attribute, final Set> examples) {
        this.attribute = attribute;

        final Set attributeValues = attribute.getValues();
        final Map>> exampleMap = new HashMap>>(attributeValues.size() + 1);

        /* Fill values */
        for (final Example example : examples) {
            final AttributeValue attributeValue = example.getValues().get(attribute);
            Set> set = exampleMap.get(attributeValue);

            if (set == null) {
                set = new HashSet>();
                exampleMap.put(attributeValue, set);
            }

            set.add(example);
        }

        final Map> localChildren = new HashMap>(attributeValues.size() + 1);
        
        /* Select next attribute for each attribute value */
        for (final Map.Entry>> entry : exampleMap.entrySet()) {
            localChildren.put(entry.getKey(), decisionNodeFactory.createNode(attribute, entry.getValue()));
        }

        this.children = localChildren;
    }

    @Override
    public Attribute getAttribute() {
        return attribute;
    }

    @Override
    public Set apply(final Item item) {
        final AttributeValue value = item.getValues().get(attribute);

        if (value == null) {
            /* If no value exists for the current attribute, use the results of all the children */
            final Set results = new HashSet();

            for (final DecisionNode node : children.values()) {
                results.addAll(node.apply(item));
            }

            return results;
        } else {
            final DecisionNode node = children.get(value);

            if (node == null) {
                return Collections.emptySet();
            } else {
                return node.apply(item);
            }
        }
    }

    @Override
    public T applySingle(final Item item) {
        final AttributeValue value = item.getValues().get(attribute);

        if (value == null) {
            /* If no value exists for the current attribute, use the results of all the children */
            T result = null;

            for (final DecisionNode node : children.values()) {
                final T tempResult = node.applySingle(item);

                if (result == null) {
                    result = tempResult;
                } else if (tempResult != null) {
                    throw new IllegalArgumentException("Ambigious result for the given item!");
                }
            }

            return result;
        } else {
            final DecisionNode node = children.get(value);

            if (node == null) {
                return null;
            } else {
                return node.applySingle(item);
            }
        }
    }
}