Caro Alisson, boa tarde!
Embora eu ainda esteja engatinhando no Python, o código me pareceu relativamente simples (depois de pronto, parece fácil mesmo! :D). Isso reforça minha convicção de que deve haver uma forma fácil de fazer isso no R também...
Vou insistir um pouco mais... e não tendo sucesso, tento adaptar o que tinha em mente para o Python!