Taro Logo

Distribute Candies to People

Easy
Amazon logo
Amazon
1 view
Topics:
Arrays

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:

  • 1 <= candies <= 10^9
  • 1 <= num_people <= 1000

Solution


Naive Approach

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.

Code (Python)

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

Time Complexity

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.

Space Complexity

O(n), where n is the number of people. We use an array of size n to store the distribution.

Optimal Approach

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.

Steps:

  1. Calculate the number of complete rounds. This can be derived by solving a quadratic equation based on the arithmetic sum. k * (1+ k) / 2, where k is the number of candies distributed during that pass and can be computed from sum(1..n), then sum(n+1...2n), etc.
  2. Determine the remaining candies.
  3. Iterate through the num_people and calculate candies for each person using the number of complete rounds.
  4. Handle remaining candies to be distributed at the end.

Code (Python)

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

Time Complexity

O(n), where n is the number of people. We iterate through the array once.

Space Complexity

O(n), where n is the number of people. We use an array of size n to store the distribution.

Edge Cases

  • candies = 0: The result should be an array of all zeros.
  • num_people = 1: The result should be an array with a single element, which is the total number of candies.
  • candies < num_people: The first few people receive candies 1, 2, 3, ..., candies, and the rest receive 0.