Suggested fix for allowing flexible labels

Xyclade · Feb 11, 2015 · 5ffaaa0 · 5ffaaa0
1 parent c80a84b
commit 5ffaaa0
Showing 1 changed file with 16 additions and 6 deletions.
diff --git a/SmilePlot/src/smile/plot/ScatterPlot.java b/SmilePlot/src/smile/plot/ScatterPlot.java
@@ -7,6 +7,8 @@
 
 import java.awt.Color;
 import java.util.Arrays;
+import java.util.HashMap;
+
 import smile.math.Math;
 
 /**
@@ -34,6 +36,7 @@ public class ScatterPlot extends Plot {
      */
     private Color[] palette;
 
+    private HashMap<Integer,Integer> valueLookupTable;
     /**
      * The legend of points.
      */
@@ -128,19 +131,26 @@ public ScatterPlot(double[][] data, int[] y, char[] legends, Color[] palette) {
         // class label set.
         int[] id = Math.unique(y);
         Arrays.sort(id);
-
+
+        valueLookupTable = new HashMap<Integer, Integer>(id.length);
+
+        for (int i =0; i <  id.length; i++) {
+            valueLookupTable.put(id[i], i);
+        }
+
         for (int i = 0; i < id.length; i++) {
             if (id[i] < 0) {
                 throw new IllegalArgumentException("Negative class label: " + id[i]); 
             }
         }
 
-        int k = Math.max(id);
-        if (legends != null && k >= legends.length) {
+        int k = id.length;
+
+        if (legends != null && k > legends.length) {
             throw new IllegalArgumentException("Too few legends.");
         }
 
-        if (palette != null && k >= palette.length) {
+        if (palette != null && k > palette.length) {
             throw new IllegalArgumentException("Too few colors.");
         }
 
@@ -167,11 +177,11 @@ public void paint(Graphics g) {
             } else {
                 for (int i = 0; i < data.length; i++) {
                     if (palette != null) {
-                        g.setColor(palette[y[i]]);
+                        g.setColor(palette[valueLookupTable.get(y[i])]);
                     }
 
                     if (legends != null) {
-                        g.drawPoint(legends[y[i]], data[i]);
+                        g.drawPoint(legends[valueLookupTable.get(y[i])], data[i]);
                     } else {
                         g.drawPoint(legend, data[i]);
                     }