基于可用DataFrame行自动调整seaborn.heatmap的形状

你的代码在数据行数变动时报错，核心原因是硬编码了固定的维度参数（比如强制reshape成3x3、手动定义长度固定的Yrows/Xcols列表），这些设定在数据量不是9的时候就会出现维度不匹配的问题。下面是修改后的代码，完全满足你的两个需求：

import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from math import ceil

# 读取数据
df = pd.read_csv("Heatmap_data.csv")

# -------------------------- 核心动态适配逻辑 --------------------------
# 可选：固定列数为3（和原逻辑一致），或者改为自动计算接近正方形的行列
# 自动计算行列（更通用）：
# num_cols = int(np.ceil(np.sqrt(len(df))))
# num_rows = int(np.ceil(len(df) / num_cols))

# 保持原逻辑的固定列数3
num_cols = 3
num_rows = ceil(len(df) / num_cols)

# 补充缺失的空行，确保数据行数刚好等于行列乘积
missing_rows = num_rows * num_cols - len(df)
if missing_rows > 0:
    empty_rows = pd.DataFrame({
        'Symbol': [''] * missing_rows,
        'Change': [np.nan] * missing_rows
    })
    df = pd.concat([df, empty_rows], ignore_index=True)

# 动态生成Yrows和Xcols，不再手动写死
Yrows = np.repeat(range(1, num_rows + 1), num_cols)
Xcols = np.tile(range(1, num_cols + 1), num_rows)
df = df.assign(Yrows=Yrows, Xcols=Xcols)

# 动态reshape，适配当前的行列数
symbol = np.asarray(df['Symbol']).reshape(num_rows, num_cols)
perchange = np.asarray(df['Change']).reshape(num_rows, num_cols)

# 生成标签，处理空值（NaN）的情况，避免格式化报错
labels = np.array([
    "{0} \n {1:.2f}%".format(symb, value) if not pd.isna(value) else "" 
    for symb, value in zip(symbol.flatten(), perchange.flatten())
]).reshape(num_rows, num_cols)

# 生成透视表
result = df.pivot(index='Yrows', columns='Xcols', values='Change')
# -------------------------- 绘图部分基本不变 --------------------------
fig, ax = plt.subplots(figsize=(13, 7))
title = "Heat Map"
plt.title(title, fontsize=18)
ttl = ax.title
ttl.set_position([0.5, 1.05])
ax.set_xticks([])
ax.set_yticks([])
ax.axis('off')

# 绘制热图：用mask隐藏空值单元格的颜色，让空单元格更美观
sns.heatmap(result, annot=labels, fmt="", cmap='RdYlGn', linewidths=0.40, ax=ax, 
            mask=result.isna(), cbar=False)

plt.show()

关键修改点说明： #

1. 动态计算行列数：
   • 你可以选择固定列数（比如保持原逻辑的3列），或者用自动计算的方式（注释掉的代码），让热图形状尽量接近正方形，适配任意数据量。
2. 自动补充空行：
   当数据行数不足行列乘积时，自动补充空行（Symbol为空，Change为NaN），确保reshape操作不会因维度不匹配报错。
3. 动态生成Yrows/Xcols：
   用np.repeat和np.tile自动生成对应的行列标识，不再依赖手动写死的列表，彻底解决长度不匹配问题。
4. 空值处理：
   生成标签时判断空值，避免对NaN进行格式化操作；绘图时用mask参数隐藏空值单元格的颜色，保持热图美观。

这样修改后，不管你的DataFrame是8行、7行还是其他任意行数，代码都能正常运行，并且自动调整热图的形状适配数据。

内容的提问来源于stack exchange，提问作者Pravat