czsem.netgraph.TreeComputation Maven / Gradle / Ivy

package czsem.netgraph;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.IntStream;

import czsem.netgraph.treesource.TreeSource;

public class TreeComputation {
	private final TreeSource treeSource;
	private final List> nodes = new ArrayList<>();
	private final List nonTreeLinks = new ArrayList<>();
	private final Map node_index = new HashMap<>();
	private List balacedOrder;
	private int maxDepth = -1;
	private List[] nodesByDepth;

	public TreeComputation(TreeSource treeSource) {
		this.treeSource = treeSource;
	}
	
	public static class NodeInfo {
		final public E node;
		final public int depth;
		final public int nodeIndex;
		final public int parentIndex;
		public int numDescendants = 0;
		final public List childrenIndexes = new ArrayList<>();

		public NodeInfo(E node, int depth, int nodeIndex, int parentIndex) {
			this.node = node;
			this.depth = depth;
			this.nodeIndex = nodeIndex;
			this.parentIndex = parentIndex;
		}

		@Override
		public String toString() {
			return node.toString();
		}
	}

	
	public void compute() {
		E root = treeSource.getRoot();
		
		if (root == null) return;
		
		addNodeAndCountDescendants(root, 0, -1);
		
		findNodesByDepth();
	}


	protected void findNodesByDepth() {
		nodesByDepth = newArray(maxDepth+1);
		for (int i = 0; i < nodesByDepth.length; i++) {
			nodesByDepth[i] = new ArrayList<>();
		}
		
		for (NodeInfo node : nodes) {
			nodesByDepth[node.depth].add(node.nodeIndex);
		}
	}


	/** fills the nodes array **/
	protected int addNodeAndCountDescendants(E node, int depth, int parentIndex) {
		if (depth > maxDepth) maxDepth = depth;
		
		
		if (node_index.containsKey(node)) {
			//existing node = "cycles" in this tree
			
			int index = node_index.get(node);
			
			
			if (parentIndex != -1) {
				//nodes.get(parentIndex).childrenIndexes.add(index);
				nonTreeLinks.add(parentIndex);
				nonTreeLinks.add(index);
			}
			
			return nodes.get(index).numDescendants;
		} 
		
		
		int index = nodes.size();
		node_index.put(node, index);
		NodeInfo info = new NodeInfo<>(node, depth, index, parentIndex);
		nodes.add(info);

		if (parentIndex != -1) {
			nodes.get(parentIndex).childrenIndexes.add(index);
		}
		
		
		int descendants = 0;
		Collection children = treeSource.getChildren(node);
		if (children != null) {
			for (E ch : children) {
				descendants += 1 + addNodeAndCountDescendants(ch, depth+1, index);
			}
		}
		
		info.numDescendants = descendants;
		
		return descendants;
	}


	public int[] collectEdges() {
		if (nodes.isEmpty()) return new int[0];
		
		int[] ret = new int[(nodes.size()-1)*2];
		
		int index = 0;
		for (NodeInfo i : nodes) {
			if (i.parentIndex == -1) continue;
			
			ret[index++] = i.parentIndex;
			ret[index++] = i.nodeIndex;
		}
		
		return ret;
	}


	public E[] collectNodes() {
		E[] ret = newArray(nodes.size());
		
		int index = 0;
		for (NodeInfo i : nodes) {
			ret[index++] = i.node; 
		}
		
		return ret;
	}
	
	
	public Integer[] computeSortedNodes() {
		Integer[] sortOredr =
				treeSource.getOrderComparator() == null
			?
				computeBalacedOrder()
			:
				computeOrder()
		;
				
		return sortOredr;
	}


	public int[] computeNodeOrder() {
		int[] ret = new int[nodes.size()];
		
		Integer[] sortOredr = computeSortedNodes();
				
		for (int r = 0; r < ret.length; r++) {
			ret[sortOredr[r]] = r;
		}

		return ret;
	}


	protected Integer[] computeOrder() {
		Comparator cmp = treeSource.getOrderComparator();
		Integer [] sortOrder = IntStream.range(0, nodes.size()).boxed().toArray(Integer[]::new);
		Arrays.sort(sortOrder, (a, b) -> cmp.compare(nodes.get(a).node, nodes.get(b).node));
		return sortOrder;
	}


	protected Integer[] computeBalacedOrder() {
		balacedOrder = new ArrayList<>(nodes.size());
		addToBalacedOrder(0);
		
		return balacedOrder.toArray(new Integer[balacedOrder.size()]);
	}


	protected void addToBalacedOrder(int nodeIndex) {
		NodeInfo nodeInfo = nodes.get(nodeIndex);
		if (nodeInfo.numDescendants == 0) {
			balacedOrder.add(nodeIndex);
			return;
		}
		
		//find best split index
		int numChildern = nodeInfo.childrenIndexes.size();
		int bestSplitIndex = 0;
		int minDiff = Integer.MAX_VALUE;
		int left = 0;
		for (int splitIndex = 0; splitIndex < numChildern; splitIndex++) {
			int diff = Math.abs(nodeInfo.numDescendants - left - splitIndex);
			if (diff < minDiff) {
				minDiff = diff;
				bestSplitIndex = splitIndex;
			}
			left += 1 + nodes.get(nodeInfo.childrenIndexes.get(splitIndex)).numDescendants;
		}

		//make the split
		for (int splitIndex = 0; splitIndex < numChildern; splitIndex++) {
			if (splitIndex == bestSplitIndex) {
				balacedOrder.add(nodeIndex);
			}
			addToBalacedOrder(nodeInfo.childrenIndexes.get(splitIndex));
		}
		
	}


	public int getDepth(int j) {
		return nodes.get(j).depth;
	}
	
	
	
	@SafeVarargs
	public static  E[] newArray(int length, E... array)
	{
	    return Arrays.copyOf(array, length);
	}


	public int getMaxDepth() {
		return maxDepth;
	}


	public int[] collectLinks() {
		return nonTreeLinks.stream().mapToInt(i->i).toArray();
	}
	
}