diff --git a/src/main/df_loader.py b/src/main/df_loader.py index 48a6942..2799f6b 100644 --- a/src/main/df_loader.py +++ b/src/main/df_loader.py @@ -2,6 +2,7 @@ from datetime import date import numpy as np import pandas as pd +from numpy import ndarray from pandas import DataFrame from src.main.constants import Constants as const @@ -73,8 +74,8 @@ class DfLoader: self.__df['location-lo'] = self.__df.loc[:, 'location'] \ .apply(lambda val: 0 if Utils.is_empty_collection(val) else val[1]) - def get_clustering_data(self) -> DataFrame: + def get_data(self) -> ndarray: columns: [] = ['location-la', 'location-lo', 'sex', 'age', 'is_university', 'is_work', 'is_student', 'is_schoolboy'] df = self.__df - return df[columns] + return df[columns].to_numpy()