Yeah, that could be problematic. The code snippet gives the general idea but it could be changed by for example by a flag telling the cpus when they enter idle to update their state_count. Or something like that.
Yes, this idea could be helpful.But since the idle path isn't a hot path. and a few memory access won't cost a lot. So I doubt if the benefit could be measurable.
But if you think the patchset is fine, it is ok, we can improve things afterwards.