diff --git a/src/ai/fitness/holes.py b/src/ai/fitness/holes.py new file mode 100644 index 0000000..5c8b110 --- /dev/null +++ b/src/ai/fitness/holes.py @@ -0,0 +1,22 @@ +from typing import Optional + +import numpy as np + +from .peaks import get_peaks + + +def get_holes( + field: np.ndarray, + peaks: Optional[np.array] = None, +) -> np.array: + if peaks is None: + peaks = get_peaks(field) + col_count = field.shape[1] + holes = np.zeros(col_count, dtype=int) + + for col in range(col_count): + start = -peaks[col] + if start != 0: + holes[col] = np.count_nonzero(field[int(start) :, col] == 0) + + return holes diff --git a/tests/ai/test_fitness.py b/tests/ai/test_fitness.py index 4f1f8f4..49cbf67 100644 --- a/tests/ai/test_fitness.py +++ b/tests/ai/test_fitness.py @@ -48,3 +48,12 @@ class TestFitness(unittest.TestCase): answers = (5, 0, 4) for field, answer in zip(self.fields, answers): self.assertEqual(get_bumpiness(None, field), answer) + + def test_get_holes(self): + answers = ( + np.array([1, 1, 0, 1, 2]), + np.array([0, 0, 0, 0, 0]), + np.array([0, 1, 0, 0, 0]), + ) + for field, answer in zip(self.fields, answers): + self.assertTrue(np.array_equal(get_holes(field), answer))