org.solovyev.common.collections.tree.DepthTreeIterator Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of common-collections-tree Show documentation

There is a newer version: 1.0.7

Show newest version

package org.solovyev.common.collections.tree;

import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Iterator;

/**
 * User: serso
 * Date: 4/1/12
 * Time: 1:36 PM
 */

/**
 * Tree iterator implementing depth-first search
 *
 * @param  type of data in node
 */
public final class DepthTreeIterator implements TreeIterator {

    @NotNull
    private Iterator> iterator;

    @Nullable
    private DepthTreeIterator childIterator;

    private final int depth;

    @Nullable
    private TreeNode lastSelfResult;

    public DepthTreeIterator(@NotNull TreeNode root) {
        this(new ArrayList>(Arrays.asList(root)));
    }

    public DepthTreeIterator(@NotNull Collection> nodes) {
        this(nodes, 0);
    }

    private DepthTreeIterator(@NotNull Collection> nodes, int depth) {
        this(nodes.iterator(), depth);
    }

    private DepthTreeIterator(@NotNull Iterator> iterator, int depth) {
        this.iterator = iterator;
        this.depth = depth;
    }

    @Override
    public boolean hasNext() {
        return iterator.hasNext() || selfHasNext() || childrenHasNext();
    }

    private boolean childrenHasNext() {
        return childIterator != null && childIterator.hasNext();
    }

    private boolean selfHasNext() {
        return lastSelfResult != null && !lastSelfResult.getOwnChildren().isEmpty();
    }

    @Override
    public TreeNode next() {
        if (selfHasNext()) {
            if (lastSelfResult instanceof MutableTreeNode) {
                childIterator = new DepthTreeIterator(((MutableTreeNode) lastSelfResult).getOwnChildrenIterator(), depth + 1);
            } else {
                childIterator = new DepthTreeIterator(lastSelfResult.getOwnChildren(), depth + 1);
            }
        }
        
        if (childrenHasNext()) {
            lastSelfResult = null;
            return this.childIterator.next();
        } else {
            childIterator = null;
            lastSelfResult = iterator.next();

            return lastSelfResult;
        }
    }

    @Override
    public int getDepth() {
        if (this.childIterator != null) {
            return this.childIterator.getDepth();
        } else {
            return this.depth;
        }
    }

    @Override
    public void remove() {
        if (this.childIterator != null) {
            this.childIterator.remove();
        } else {
            this.iterator.remove();
        }
    }
}