com.blazebit.ai.decisiontree.impl.StringDecisionNode 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 com.blazebit.collection.TrieMap;

import java.util.Collections;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

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

    private final Attribute attribute;
    private final TrieMap> children;

    public StringDecisionNode(final DecisionNodeFactory decisionNodeFactory, final Attribute attribute, final Set> examples) {
        this.attribute = attribute;
        this.children = new TrieMap>();

        final TrieMap>> exampleTrieMap = new TrieMap>>();

        /* Fill values */
        for (final Example example : examples) {
            final AttributeValue exampleAttributeValue = example.getValues().get(attribute);
            final String key = exampleAttributeValue == null ? "" : (String) exampleAttributeValue.getValue();
            Set> set = exampleTrieMap.get(key);

            if (set == null) {
                set = new HashSet>();
                exampleTrieMap.put(key, set);
            }

            set.add(example);
        }

        final TrieMap> localChildren = children;
        
        /* Select next attribute for each attribute value */
        for (final Map.Entry>> entry : exampleTrieMap.entrySet()) {
            final Set> resultExamples = entry.getValue();

            if (resultExamples.size() > 0) {
                localChildren.put(entry.getKey(), decisionNodeFactory.createNode(attribute, resultExamples));
            }
        }
    }

    @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((String) value.getValue());

            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);
            }
        }
    }
}