How to add legend/gradient in Folium Heat Map?

Question

I am creating a heat map using Folium.My data contains 3 columns one is category, lat and long. The lat-long points are categorized into 3 categories like A, B, and C. I am able to plot the heat map using folium, but I need to add the legend showing the color difference between the points.I need to mark points into 3 different colors based on the category.

I am attaching the sample code which for your reference.Any help is appreciated.

Thanks in advance!

from folium import plugins
from folium.plugins import HeatMap
from folium.plugins import MarkerCluster
import pandas as pd

map = folium.Map(location=[lat, long],zoom_start =12) 
data = pd.read_csv(filename)
# List comprehension to make out list of lists
heat_data = [[row['LAT'],row['LONG'],] for index, row in data.iterrows()]
# Plot it on the map
HeatMap(heat_data).add_to(map)
# Display the map
map
map.save('C:\Temp\map2.html')

Answer 1

Below is the solution and it's based on the answer of @Alexei Novikov Code below is more complete

import branca.colormap
from collections import defaultdict
import folium
import webbrowser
from folium.plugins import HeatMap 

map_osm = folium.Map(llocation=[35,110],zoom_start=1)

steps=20
colormap = branca.colormap.linear.YlOrRd_09.scale(0, 1).to_step(steps)
gradient_map=defaultdict(dict)
for i in range(steps):
    gradient_map[1/steps*i] = colormap.rgb_hex_str(1/steps*i)
colormap.add_to(map_osm) #add color bar at the top of the map

HeatMap(data1,gradient = gradient_map).add_to(map_osm) # Add heat map to the previously created map

file_path = r"C:\\test.html"
map_osm.save(file_path) # Save as html file
webbrowser.open(file_path) # Default browser open