Fix bug in updating the state based on a mask, in update_with_mask function

wordle_env/state.py

@@ -77,17 +77,12 @@ def update_from_mask(state: WordleState, word: str, mask: List[int]) -> WordleSt
             _set_no(state, offset, i)
             _set_if_cero(state, offset, [0, 1, 0])
         elif mask[i] == NO:
-            # Need to check this first in case there's prior maybe + yes
-            if c in prior_maybe:
-                # Then the maybe could be anywhere except here
-                state[offset + 3 * i : offset + 3 * i + 3] = [1, 0, 0]
-            elif c in prior_yes:
-                # No maybe, definitely a yes,
-                # so it's zero everywhere except the yesses
-                for j in range(WORDLE_N):
-                    # Only flip no if previously was maybe
-                    if state[offset + 3 * j : offset + 3 * j + 3][1] == 1:
-                        state[offset + 3 * j : offset + 3 * j + 3] = [1, 0, 0]
+            # Need to check this first in case there's prior maybe or yes
+            if c in prior_yes or c in prior_maybe:
+                # Definitely not here
+                _set_no(state, offset, i)
+                # It's zero everywhere except the yesses and maybes
+                _set_if_cero(state, offset, [1, 0, 0])
             else:
                 # Just straight up no
                 _set_all_no(state, offset)