如何在matplotlib网格的特定单元格中着色？

import matplotlib.pyplot as plt
import numpy as np

# 创建一个2x2的网格
grid = np.zeros((2, 2))

# 在特定单元格中设置值
grid[0, 0] = 1
grid[1, 1] = 2

# 创建一个图形对象和子图对象
fig, ax = plt.subplots()

# 使用imshow函数在网格中着色
ax.imshow(grid, cmap='viridis')

# 隐藏坐标轴
ax.axis('off')

# 显示图形
plt.show()