We distribute some number of candies
to a row of n = num_people
people in the following way:
We then give 1 candy to the first person, 2 candies to the second person, and so on until we give n
candies to the last person.
Then, we go back to the start of the row, giving n + 1
candies to the first person, n + 2
candies to the second person, and so on until we give 2 * n
candies to the last person.
This process repeats (with us giving one more candy each time, and moving to the start of the row after we reach the end) until we run out of candies. The last person will receive all of our remaining candies (not necessarily one more than the previous gift).
Return an array (of length num_people
and sum candies
) that represents the final distribution of candies.
Example 1:
Input: candies = 7, num_people = 4
Output: [1,2,3,1]
Explanation:
On the first turn, ans[0] += 1, and the array is [1,0,0,0].
On the second turn, ans[1] += 2, and the array is [1,2,0,0].
On the third turn, ans[2] += 3, and the array is [1,2,3,0].
On the fourth turn, ans[3] += 1 (because there is only one candy left), and the final array is [1,2,3,1].
Example 2:
Input: candies = 10, num_people = 3
Output: [5,2,3]
Explanation:
On the first turn, ans[0] += 1, and the array is [1,0,0].
On the second turn, ans[1] += 2, and the array is [1,2,0].
On the third turn, ans[2] += 3, and the array is [1,2,3].
On the fourth turn, ans[0] += 4, and the final array is [5,2,3].
Constraints:
The most straightforward approach is to simulate the candy distribution process as described in the problem. We iterate through the people, giving them an increasing number of candies until we run out. This brute force solution is easy to understand and implement.
def distribute_candies_naive(candies: int, num_people: int) -> list[int]:
result = [0] * num_people
give = 1
i = 0
while candies > 0:
result[i % num_people] += min(candies, give)
candies -= give
give += 1
i += 1
return result
O(k), where k is the number of candies distributed. In the worst case, where we have a very large number of candies, this could be quite high.
O(n), where n is the number of people. We use an array of size n
to store the distribution.
The optimal solution involves some mathematical insights to avoid the iterative candy distribution. First, determine how many complete rounds of distribution occur. A complete round means each person receives candy once.
Next, we can calculate the total number of candies distributed in full rounds using the sum of arithmetic series. Then subtract these candies from the initial candy count. Finally, distribute the remaining candies.
num_people
and calculate candies for each person using the number of complete rounds.import math
def distribute_candies(candies: int, num_people: int) -> list[int]:
n = num_people
p = int((math.sqrt(2 * candies + 0.25) - 0.5) / n)
remaining_candies = candies - (n * p * (p + 1)) // 2
ans = [0] * n
for i in range(n):
ans[i] = (p * (i + 1)) + (n * (p * (p - 1)) // 2 if p > 0 else 0)
current = i + 1 + p * n
if remaining_candies > 0:
if current <= remaining_candies:
ans[i] += current
remaining_candies -= current
else:
ans[i]+= remaining_candies
remaining_candies = 0
return ans
O(n), where n is the number of people. We iterate through the array once.
O(n), where n is the number of people. We use an array of size n
to store the distribution.