Fixed bug in partitioning function