44 lines
1.1 KiB
Python
44 lines
1.1 KiB
Python
from pandas import DataFrame, read_csv
|
|
|
|
|
|
def replace_values(df) -> DataFrame:
|
|
for column in df.columns:
|
|
df[column].fillna(value=df[column].mean(), inplace=True)
|
|
return df
|
|
|
|
|
|
def process_na(df, action) -> DataFrame:
|
|
if action == "drop":
|
|
return df.dropna()
|
|
elif action == "fill":
|
|
return replace_values(df)
|
|
else:
|
|
print("Unknown action selected. The choices are: ")
|
|
print("fill: fills the na values with the mean")
|
|
print("drop: drops the na values")
|
|
exit()
|
|
|
|
|
|
def filter_dataframe(df) -> DataFrame:
|
|
relevant_columns = [
|
|
"HORA",
|
|
"DIASEMANA",
|
|
"COMUNIDAD_AUTONOMA",
|
|
"ISLA",
|
|
"TOT_HERIDOS_LEVES",
|
|
"TOT_HERIDOS_GRAVES",
|
|
"TOT_VEHICULOS_IMPLICADOS",
|
|
"TOT_MUERTOS",
|
|
"TIPO_VIA",
|
|
"LUMINOSIDAD",
|
|
"FACTORES_ATMOSFERICOS",
|
|
]
|
|
filtered_df = df.filter(items=relevant_columns)
|
|
return filtered_df
|
|
|
|
|
|
def parse_data(source, action) -> DataFrame:
|
|
df = read_csv(filepath_or_buffer=source, na_values="?")
|
|
processed_df = process_na(df=df, action=action)
|
|
return processed_df
|