プリム法による最小全域木を求めるアレをなんとか自力で実装してみた

Prim.java

class Prim<T>
{
        private Set<T> rest;
	private List<Path<T>> paths;
	private Set<T> already;

	public Prim()
	{
		rest = new HashSet<T>();
		paths = new LinkedList<Path<T>>();
		already = new HashSet<T>();
	}
	public void addNode(T node)
	{
		rest.add(node);
	}
	public void addPath(T a, T b, int cost)
	{
		paths.add(new Path<T>(a, b, cost));
	}
	public Set<Path<T>> getResult()
	{
		Collections.sort(paths);
		Set<Path<T>> result = new HashSet<Path<T>>();
		{
			Iterator<T> itr = rest.iterator();
			already.add(itr.next());
			itr.remove();
		}
		while(!rest.isEmpty())
		{
			Iterator<Path<T>> itr = paths.iterator();
			while(itr.hasNext())
			{
				Path<T> p = itr.next();
				ConnectionStatus status = p.canConnect(already);
				if(status==ConnectionStatus.REMOVABLE) itr.remove();
				else if(status==ConnectionStatus.NONE);
				else
				{
					T connect = p.getConnectable(status);
					if(!rest.remove(connect)) throw new IllegalStateException();
					already.add(connect);
					result.add(p);
					itr.remove();
					break;
				}
			}
		}
		return result;
	}

	static class Path<T> implements Comparable<Path<T>>
	{
		private final T a, b;
		private final int cost;

		Path(T a, T b, int c)
		{
			this.a = a;
			this.b = b;
			this.cost = c;
		}
		ConnectionStatus canConnect(Set<T> s)
		{
			if(s.contains(a))
			{
				if(s.contains(b))
					return ConnectionStatus.REMOVABLE;
				else return ConnectionStatus.A;
			}
			else
			{
				if(s.contains(b))
					return ConnectionStatus.B;
				else return ConnectionStatus.NONE;
			}
		}
		T getConnectable(ConnectionStatus status)
		{
			switch(status)
			{
				case A: return b;
				case B: return a;
				default: return null;
			}
		}
		@Override
		public int compareTo(Path<T> o)
		{
			return cost>o.cost ? 1 : cost==o.cost ? 0 : -1;
		}
		public int getCost()
		{
			return cost;
		}
		public T getA()
		{
			return a;
		}
		public T getB()
		{
			return b;
		}
		@Override
		public String toString()
		{
			return a + " - " + b + " : " + cost;
		}
	}

	public static enum ConnectionStatus {REMOVABLE, A, B, NONE}
}