diff --git a/mlrose/fitness.py b/mlrose/fitness.py
index 39d8b8ff..a17ddce7 100644
--- a/mlrose/fitness.py
+++ b/mlrose/fitness.py
@@ -892,11 +892,11 @@ class MaxKColor:
 
         >>> import mlrose
         >>> import numpy as np
-        >>> edges = [(0, 1), (0, 2), (0, 4), (1, 3), (2, 0), (2, 3), (3, 4)]
+        >>> edges = [(0, 1), (0, 2), (0, 4), (1, 3), (2, 0), (2, 3), (3, 4), (0, 5)]
         >>> fitness = mlrose.MaxKColor(edges)
-        >>> state = np.array([0, 1, 0, 1, 1])
+        >>> state = np.array([0, 1, 0, 1, 1, 0])
         >>> fitness.evaluate(state)
-        3
+        4
 
     Note
     ----
diff --git a/tests/test_fitness.py b/tests/test_fitness.py
index f7aeb47a..04739273 100644
--- a/tests/test_fitness.py
+++ b/tests/test_fitness.py
@@ -183,10 +183,10 @@ def test_queens():
     @staticmethod
     def test_max_k_color():
         """Test MaxKColor fitness function"""
-        edges = [(0, 1), (0, 2), (0, 4), (1, 3), (2, 0), (2, 3), (3, 4)]
+        edges = [(0, 1), (0, 2), (0, 4), (1, 3), (2, 0), (2, 3), (3, 4), (0, 5)]
 
-        state = np.array([0, 1, 0, 1, 1])
-        assert MaxKColor(edges).evaluate(state) == 3
+        state = np.array([0, 1, 0, 1, 1, 0])
+        assert MaxKColor(edges).evaluate(state) == 4
 
     @staticmethod
     def test_custom_fitness():