diff --git a/prism/src/explicit/StateValues.java b/prism/src/explicit/StateValues.java index f0566de1..a34fe284 100644 --- a/prism/src/explicit/StateValues.java +++ b/prism/src/explicit/StateValues.java @@ -32,6 +32,8 @@ import java.io.FileReader; import java.io.IOException; import java.util.BitSet; import java.util.List; +import java.util.Map; +import java.util.TreeMap; import parser.State; import parser.ast.ExpressionBinaryOp; @@ -1724,6 +1726,48 @@ public class StateValues implements StateVector } } + /** Partition the state values by (boolean) value, returns a Map: Value -> States-with-that-value */ + public Map partitionBool() { + if (!type.equals(TypeBool.getInstance())) + throw new UnsupportedOperationException("Can not invoke paritionBool() for StateValues with "+type+" elements."); + return this.partition(); + } + + /** Partition the state values by (double) value, returns a Map: Value -> States-with-that-value */ + public Map partitionDouble() { + if (!type.equals(TypeDouble.getInstance())) + throw new UnsupportedOperationException("Can not invoke paritionDouble() for StateValues with "+type+" elements."); + return this.partition(); + } + + /** Partition the state values by (integer) value, returns a Map: Value -> States-with-that-value */ + public Map partitionInt() { + if (!type.equals(TypeInt.getInstance())) + throw new UnsupportedOperationException("Can not invoke paritionInt() for StateValues with "+type+" elements."); + return this.partition(); + } + + @SuppressWarnings("unchecked") + /** + * Helper: Partition the state values by value. + * Returns a Map: Value -> States-with-that-value + * @param the type of the values + */ + private > Map partition() { + Map result = new TreeMap(); + + for (int i = 0; i < getSize(); i++) { + BitSet set = result.get(getValue(i)); + if (set == null) { + set = new BitSet(); + result.put((T) getValue(i), set); + } + set.set(i); + } + + return result; + } + /** * Make a (deep) copy of this vector */