Program to find hoe many children will get candies while distributing them maintaining the rules in Python

Suppose we have k number of candies. We have to distribute them among children. Now there are some rules

• ith child will get i^2 number of candies
• any children at index i will not get any candy until all children from index 1 to i-i are served
• If ith children does not get i^2 number of candies, then that is not a valid serve.

So, if the input is like k = 20, then the output will be 3 because, first one will get 1, second one will get 2^2 = 4, third one will get 3^2 = 9, but fourth one needs 4^2 = 16, but we have only 6 candies left, so this is not a valid distribution, so only three children will be served.

To solve this, we will follow these steps −

• left := 0, right := k
• while right - left > 1, do
  • mid := floor of (left + right) / 2
  • if floor of (mid * (mid + 1) * (2 * mid + 1) / 6) > k, then
    • right := mid
  • otherwise,
    • left := mid
• if right *(right + 1) *(2 * right + 1) <= k * 6, then
  • return right
• return left

Example

Let us see the following implementation to get better understanding −

def solve(k):
   left = 0
   right = k
   while (right - left > 1):
      mid = (left + right) // 2
      if (mid * (mid + 1) * (2 * mid + 1) // 6 > k):
         right = mid
      else:
         left = mid
   if (right * (right + 1) * (2 * right + 1) <= k * 6):
      return right
   return left

k = 20
print(solve(k))

Input

20

Output

3