In [3]:
# Set attributes in relevant format for k-means clustering
df['release_year'] = pd.to_datetime(df['original_release_date']).dt.year
def encodeAttributes(col, n):
# Use one-hot encoding for top N entries in column attribute
top_genres = df[col].value_counts().index[:n]
return top_genres
def appendAttributes(col, n):
# Append encoded attributes to original dataframe
for item in encodeAttributes(col, n):
df[item] = df[col].fillna('').str.lower().apply(lambda x: 1 if item.lower() in x else 0)
return df
# Select most relevant features - manually selected for time-being
df = appendAttributes('genres', 10)
df = appendAttributes('directors', 10)
selected_attributes = ['tomatometer_rating', 'release_year', 'audience_count'] + list(encodeAttributes('genres', 10)) + list(encodeAttributes('directors', 10))
# Standardise feature values and impute missing values
scaler = StandardScaler()
imputer = SimpleImputer(strategy='mean')
X = scaler.fit_transform(df[selected_attributes])
X = imputer.fit_transform(X)
# Plot heatmap showing correlation matrix of selected attributes
plt.figure(figsize=(8, 6))
heatmap = plt.imshow(df[selected_attributes].corr(), cmap='coolwarm', interpolation='nearest')
plt.title("Attribute Correlation Heatmap")
plt.xticks(np.arange(df[selected_attributes].corr().shape[1]), labels=df[selected_attributes].corr().columns, rotation=90)
plt.yticks(np.arange(df[selected_attributes].corr().shape[1]), labels=df[selected_attributes].corr().columns)
plt.colorbar(heatmap, fraction=0.046, pad=0.04)
plt.tight_layout()
plt.show()
