[BUG] Corrected a bug in the random walker that appeared

* when returning the full probability instead of the segmentation
    * for three or more labels (two was OK)

skimage/segmentation/random_walker_segmentation.py

@@ -401,9 +401,8 @@ def random_walker(data, labels, beta=130, mode='bf', tol=1.e-3, copy=True,
                      copy=True).reshape(dims) for Xline in X])
         for i in range(1, int(labels.max()) + 1):
             mask_i = np.squeeze(labels == i)
+            X[:, mask_i] = 0
             X[i - 1, mask_i] = 1
-            X[np.setdiff1d(np.arange(0, labels.max(), dtype=np.int),
-                           [i - 1]), mask_i] = 0
     else:
         X = _clean_labels_ar(X + 1, labels).reshape(dims)
     return X

skimage/segmentation/tests/test_random_walker.py

@@ -58,8 +58,13 @@ def test_2d_bf():
                                     return_full_prob=True)
     assert (full_prob_bf[1, 25:45, 40:60] >=
                         full_prob_bf[0, 25:45, 40:60]).all()
-    return data, labels_bf, full_prob_bf
-
+    # Now test with more than two labels
+    labels[55, 80] = 3
+    full_prob_bf = random_walker(data, labels, beta=90, mode='bf',
+                                    return_full_prob=True)
+    assert (full_prob_bf[1, 25:45, 40:60] >=
+                        full_prob_bf[0, 25:45, 40:60]).all()
+    assert len(full_prob_bf) == 3
 
 def test_2d_cg():
     lx = 70