feat(ai): add aggregate height calculation

This commit is contained in:
Kristofers Solo
2024-01-06 17:56:02 +02:00
parent e265581262
commit 41cef03f50
2 changed files with 36 additions and 0 deletions

16
src/ai/moves/height.py Normal file
View File

@@ -0,0 +1,16 @@
from typing import Any, Optional
import numpy as np
def aggregate_height(field: np.ndarray[int, Any]) -> int:
"""
Calculates the aggregate height of the field.
Args:
field: 2D array representing the game field.
Returns:
The aggregate height of the field.
"""
return np.sum(field.shape[0] - np.argmax(field, axis=0))

20
tests/ai/test_moves.py Normal file
View File

@@ -0,0 +1,20 @@
import unittest
import numpy as np
from ai.moves.height import aggregate_height
class TestFitness(unittest.TestCase):
def test_aggregate_height(self):
field = np.array(
[
[0, 0, 0, 0, 1, 1, 0, 0, 0, 0],
[0, 1, 1, 1, 1, 1, 1, 0, 0, 1],
[0, 1, 1, 0, 1, 1, 1, 1, 1, 1],
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
[1, 1, 1, 0, 1, 1, 1, 1, 1, 1],
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
]
)
self.assertEqual(aggregate_height(field), 48)