Fixed bug in naive solution when K > N.
Split and moved functions for clarity.
Added comments on integer overflow.
Removed trailing whitespace.