Candy distribution with two passes in Python

Why this matters

A simpler problem gives each child one cookie. This problem is different. Each child is compared with a neighbour on the left and a neighbour on the right. One sweep across the row cannot see both neighbours, so the answer needs two sweeps. A sweep is a loop that walks through a list in one direction, one position at a time.

Any rule that compares an item with both of its neighbours has this shape. Smoothing a series of temperature readings, where each value depends on both neighbours, needs a left pass and a right pass. This lesson uses candy because the rule is easy to picture, and because a single pass gives a wrong answer that is easy to see.

The problem in plain words

Children stand in a line. Each child has a rating, which is a whole number. A teacher wants to hand out candy under two rules:

1. Every child gets at least one candy.
2. A child with a higher rating than a neighbour gets more candy than that neighbour.

The goal is the smallest total number of candies that follows both rules.

Here is a small example. The ratings are 1, 0 and 2, in line order. Store them in a Python list with ratings = [1, 0, 2]. A list is an ordered collection of values written in square brackets. Each value has a position called an index, and the first index is 0. So ratings[0] is 1, ratings[1] is 0 and ratings[2] is 2.

Try the candy counts [1, 1, 2]. Child 0 has rating 1, which is higher than child 1's rating of 0. So child 0 must get more candy than child 1. Here both get 1, so the rule breaks. Try [2, 1, 2] instead. Child 0 gets 2, which is more than child 1's 1. Child 2 has rating 2, higher than child 1's rating of 0, and gets 2, which is more than 1. Both rules hold, and the total is 5.

Equal ratings put no rule between two children. Child 1 and child 2 in [1, 2, 2] have the same rating of 2, so neither needs more than the other. The next section uses this case.

Why one pass is not enough

Start with the obvious idea. Walk from left to right. For each child, if their rating is higher than the child on their left, give them one more candy than that child. Otherwise give them 1.

This handles every case where a child is higher than the child on the left. It does not handle a child who is higher than the child on the right. Take the ratings 1, 3, 2, 2, 1. The left pass gives the counts 1, 2, 1, 1, 1. Look at the last two children. The child with rating 2 has a neighbour with rating 1 on the right, so that child must get more candy. Both have 1 candy, so the rule breaks.

A right-to-left pass alone fails in the other direction. Take the ratings 1, 2, 3, 1. The right pass gives [1, 1, 2, 1]. The child with rating 2 has the same count as the child with rating 1 on its left, which breaks the rule.

The program below runs both single passes and checks each result with is_valid. That function looks at every pair of neighbours and returns False if one pair breaks a rule. The code uses range, which makes a sequence of index numbers. range(1, len(ratings)) gives the numbers 1, 2, and so on up to the last index. len returns how many items a list has.

The output is:

Left pass only: [1, 2, 1, 1, 1] valid=False
Right only on [1, 2, 3, 1] -> [1, 1, 2, 1] valid=False

The line [1] * len(ratings) makes a list with one 1 for each child. A list times a number repeats the list, so [1] * 3 is [1, 1, 1]. The text valid=False is how Python prints the value False. Python writes the two boolean values as True and False, with capital letters.

So you need both sweeps.

Two sweeps and the larger number

The method uses two sweeps, and each child keeps the larger of their two numbers. The function max(a, b) returns the larger of two numbers.

Step 1, left to right. Start with 1 candy for everyone. If the child's rating is higher than the left neighbour's rating, set the child's count to the left neighbour's count plus 1. This enforces the left rule.

Step 2, right to left. If the child's rating is higher than the right neighbour's rating, the child's count must be more than the right neighbour's count. Set the child's count to the larger of its current value and the right neighbour's count plus 1. This enforces the right rule.

Why take the larger value in step 2? After step 1, the child's count already satisfies the left rule, so it may be too small for the right rule. The larger value satisfies both. It never lowers a value, so the left rule still holds.

The reason the result is the smallest total is this. Any valid answer must give a child at least as many candies as the run of rising ratings that ends at that child from the left. The left sweep computes exactly that number. The same holds for the right side. Taking the larger of the two is the smallest count that meets both.

Think of the left sweep's counts as a staircase. Each step up happens where a rating rises from the child before. Ratings 1, 2, 3 give the counts 1, 2, 3. When a rating falls or stays level, the staircase drops back to 1. The right sweep builds the same kind of staircase from the other end. A child at the top of a staircase in either direction needs the taller of the two. Equal ratings matter in this picture because a level step breaks the staircase instead of adding to it.

Here is the code. The first loop is the left sweep. It starts at index 1, because child 0 has no left neighbour. The second loop is the right sweep. It starts at n - 2, because the last child has no right neighbour. The third number in range(n - 2, -1, -1) is the step, which is -1, so the loop counts down. The stop value -1 is not included, so the loop ends at index 0.

The function has type hints, list[int], which say that the input and the output are lists of whole numbers. Python does not check type hints when the program runs. They help a reader see what goes in and out.

The line if __name__ == "__main__": means the loop at the bottom runs only when you start this file directly. If another file imports candies, the test loop does not run.

The f"..." text is an f-string. Values written inside curly braces in an f-string are replaced by their values when the line runs. sum(c) adds up all the numbers in the list c.

The output is:

[1, 0, 2] -> [2, 1, 2] total 5
[1, 2, 2] -> [1, 2, 1] total 4
[1, 3, 2, 2, 1] -> [1, 2, 1, 2, 1] total 7
[1, 2, 87, 87, 87, 2, 1] -> [1, 2, 3, 1, 3, 2, 1] total 13

Readers who know Java will see that a Python list plays the role of an int array here. len(ratings) replaces ratings.length, and max(a, b) replaces Math.max.

The cost of the method is simple to count. Each sweep visits every child once, so the two sweeps take about twice as many steps as there are children. The extra list holds one number per child. Computer scientists write this as O(n) time and O(n) extra space, where n is the number of children. The letter O in that notation means the work grows in proportion to the input, so doubling the row roughly doubles the steps.

Worked example 1: a valley

Ratings 1, 0, 2. The left sweep starts with [1, 1, 1]. Index 1 has rating 0, which is not higher than rating 1 on its left, so it stays 1. Index 2 has rating 2, which is higher than 0, so it becomes candy[1] + 1 = 2. The left result is [1, 1, 2].

The right sweep looks at index 1 first. Its rating 0 is not higher than the rating 2 on its right, so nothing changes. Index 0 has rating 1, which is higher than 0. Its count becomes the larger of its current 1 and candy[1] + 1 = 2. The final result is [2, 1, 2], with total 5.

Worked example 2: equal neighbours

Ratings 1, 2, 2. The left sweep gives index 1 the count 2, because 2 is higher than 1. Index 2 has rating 2, which is equal to the rating before it, so it stays 1. The left result is [1, 2, 1].

The right sweep finds nothing to change. Index 1 has rating 2, which is not higher than the rating 2 on its right. Index 0 has rating 1, which is not higher than rating 2. The final result is [1, 2, 1], with total 4. Equal ratings never force extra candy, so child 2 is allowed to have less.

Worked example 3: a peak in the middle

Ratings 1, 3, 2, 2, 1. The left sweep gives [1, 2, 1, 1, 1], which was shown to be invalid above. The right sweep fixes it:

• Index 3 has rating 2 and its right neighbour has rating 1. Its count becomes max(1, 1 + 1) = 2.
• Index 2 has rating 2 and its right neighbour at index 3 also has rating 2. No change.
• Index 1 has rating 3 and its right neighbour has rating 2. Its count becomes max(2, candy[2] + 1) = max(2, 2) = 2.
• Index 0 has rating 1 and its right neighbour has rating 3. No change.

The final result is [1, 2, 1, 2, 1], with total 7.

Worked example 4: a plateau

Ratings 1, 2, 87, 87, 87, 2, 1. This is longer, and the three children with rating 87 form a plateau. Their ratings are equal, so no rule applies between them.

The left sweep gives [1, 2, 3, 1, 1, 1, 1]. Index 2 has the count 3 because it is part of the rising run 1, 2, 87. Index 3 has the count 1 because its rating equals index 2's rating, so the left sweep does not raise it.

The right sweep raises index 5 to 2, then index 4 to 3, because index 4 has rating 87 above index 5's rating 2. Indexes 3 and 2 stay as they are. The final result is [1, 2, 3, 1, 3, 2, 1], with total 13.

Checking the answer with a brute force search

A brute force search tries every possible count for every child and keeps the smallest total that follows both rules. It is slow, but it is a fair judge on short rows. A minimal answer never gives a child more candy than the number of children in the row, so the search tries counts from 1 up to that number.

The search uses itertools.product. itertools is a standard library module with tools for looping over combinations. product(range(1, n + 1), repeat=n) produces every group of n counts, where each count is between 1 and n. Each group is a tuple, which is a fixed group of values written in round brackets, such as (2, 1, 2). A tuple cannot be changed after it is made.

The list comprehension in brute_min collects the total of every valid group. A list comprehension is a short way to build a list with a loop inside square brackets. The if at the end keeps only the groups that pass is_valid.

This code uses candies and is_valid from the sections above.

from itertools import product
 
 
def brute_min(ratings: list[int]) -> int:
    n = len(ratings)
    totals = [
        sum(candy)
        for candy in product(range(1, n + 1), repeat=n)
        if is_valid(ratings, candy)
    ]
    return min(totals)

The second test uses random rows. random.seed(1) fixes the random numbers, so every run makes the same rows.

import random
 
print("Brute force on [1, 3, 2, 2, 1]:", brute_min([1, 3, 2, 2, 1]))
 
random.seed(1)
checked = 0
mismatches = 0
invalid = 0
for _ in range(400):
    n = random.randint(1, 5)
    ratings = [random.randint(0, 2) for _ in range(n)]
    fast = candies(ratings)
    checked += 1
    if not is_valid(ratings, fast):
        invalid += 1
    if sum(fast) != brute_min(ratings):
        mismatches += 1
print(f"random checks: {checked}, mismatches: {mismatches}, invalid: {invalid}")

The output was:

Brute force on [1, 3, 2, 2, 1]: 7
random checks: 400, mismatches: 0, invalid: 0

This is evidence and not a proof. The two sweeps are justified by the staircase argument. The test catches wrong implementations on small rows like these.

Common mistakes

Using only one pass

A single left pass gives a valid answer only when every rule points left. The output of the left-only run above shows this. The fix is the second sweep.

Overwriting in the second sweep

A tempting mistake is to write candy[i] = candy[i + 1] + 1 in the second sweep, without taking the maximum. This throws away the left sweep's work. Take the ratings 1, 2, 3, 1. The left sweep gives [1, 2, 3, 1]. Overwriting in the right sweep turns index 2 into candy[3] + 1 = 2, which is no longer more than index 1's 2.

Copy candies, change its right-sweep line to candy[i] = candy[i + 1] + 1, and run both versions. The output is:

Overwrite on [1, 2, 3, 1] gives [1, 2, 2, 1] total 6 valid=False
Max version on [1, 2, 3, 1] gives [1, 2, 3, 1] total 7 valid=True

The overwrite version fails the validity check, and the maximum version passes.

Treating equal ratings as a rule

The rule only applies when two neighbours differ. A child with the same rating as a neighbour needs no extra candy. The symbol >= means greater than or equal to. If both loops in candies use >= instead of >, equal neighbours are treated as if one were higher, and the total grows. Take the ratings 1, 2, 2. The correct answer is [1, 2, 1], with total 4. With >= in both loops, the output is:

At-least version on [1, 2, 2] -> [1, 4, 3] total 8 valid=True
Correct version on [1, 2, 2] -> [1, 2, 1] total 4

The >= answer still passes the validity check, because that check ignores equal neighbours. It is still wrong, because it spends 4 extra candies. This mistake does not crash the program. It only shows up when you compare the total with the true minimum.

Starting the right sweep at the last index

The right sweep reads ratings[i + 1]. For the last child, i + 1 is one past the end of the list. Starting at n - 1 instead of n - 2 causes this error on the ratings 1, 2:

IndexError: list index out of range

Python prints a traceback first. A traceback is the list of lines that ran before the error, with the line number where it stopped. The last line is the error. Python does not name the bad index in this message, so look at the line number in the traceback to find the cause. The loop must start at n - 2.

Skipping the step in range

A quieter mistake is to leave out the step in the right sweep, and write range(n - 2, -1). With no step, range counts up by 1. A range that starts above its stop value is empty, so the loop body never runs. No error appears, and the answer is wrong:

Missing step on [1, 3, 2, 2, 1]: [1, 2, 1, 1, 1] valid=False

The output is the left-only result, and the validity check catches it.

Starting every count at 0

Each child must get at least one candy. The code fills the list with 1 first, using [1] * n. If the list starts at 0, a child with no rising neighbours keeps 0 candies, which breaks rule 1. Changing candy = [1] * n to candy = [0] * n gives this result for the ratings 5, 5, 5:

Start at 0 on [5, 5, 5]: [0, 0, 0]

Practice

Exercise 1. Ratings [1, 2, 3, 4]. What are the candy counts and the total?

Show the answer

The counts are [1, 2, 3, 4] and the total is 10. The left sweep gives those values, and the right sweep changes nothing because the ratings never fall from left to right.

Exercise 2. Ratings [4, 3, 2, 1]. What are the candy counts and the total?

Show the answer

The counts are [4, 3, 2, 1] and the total is 10. The left sweep gives all 1s. The right sweep raises index 2 to 2, then index 1 to 3, then index 0 to 4.

Exercise 3. Ratings [5, 5, 5]. What are the candy counts and the total?

Show the answer

The counts are [1, 1, 1] and the total is 3. Equal ratings put no rule between neighbours, so every child gets the minimum.

Exercise 4. Ratings [1, 3, 2]. What are the candy counts and the total?

Show the answer

The counts are [1, 2, 1] and the total is 4. The left sweep gives [1, 2, 1]. In the right sweep, index 1 has rating 3, which is above rating 2 on its right. Its count becomes max(2, 1 + 1) = 2, so nothing changes.

Exercise 5. Ratings [1, 0, 0, 2]. What are the candy counts and the total?

Show the answer

The counts are [2, 1, 1, 2] and the total is 6. The left sweep gives [1, 1, 1, 2]. The right sweep changes only index 0, because its rating 1 is above the rating 0 on its right. Its count becomes max(1, 1 + 1) = 2.

Exercise 6. Ratings [1, 2, 2], but both loops use >=. What are the counts and the total, and why is the total larger than the correct one?

Show the answer

The counts are [1, 4, 3] and the total is 8. The >= test treats the equal rating of child 2 as higher than the rating of child 1, so child 2 must get more candy than child 1. The correct total is 4, with counts [1, 2, 1].

Summary

• A rule that compares a child with both neighbours needs two sweeps, one left to right and one right to left.
• The left sweep enforces the rule against the left neighbour, and the right sweep enforces it against the right neighbour.
• Keep the larger of the two numbers at each child. Overwriting in the second sweep can break the first sweep's rule.
• Both sweeps read each child once, so the work grows in a straight line with the number of children.