Learn to Code by Solving Problems: A Python Programming Primer Homepage for the book Learn to Code by Solving Problems: A Python Programming Primer

Solving USACO Photoshoot in Two Different Ways

Hello everyone!

I’ve been getting some reader requests to do another example of a complete search algorithm like in Chapter 9.

So here we go :) I’m assuming that you’ve already studied the first nine chapters of Learn to Code by Solving Problems. But I’m going to try my best to make this post useful for anyone who has the required background Python knowledge.

The Problem

The problem we’ll solve is called Photoshoot. It’s a problem from the USACO January 2020 Bronze competition.

Please read the original problem description and then keep reading here.

This problem didn’t quite make the cut for inclusion in Chapter 9. I can’t exactly remember why, though :D One reason I like this problem is because there are both $O(n^2)$ and $O(n)$ solutions. The first one, the $O(n^2)$ one, is what we’re going to work on today. It’ll be fast enough to pass all test cases in time, and it’s a nice example of a Complete Search Algorithm. … But maybe you’ll end up being super curious how we can solve this thing faster, in only $O(n)$ time. Well, if you are, stay tuned, because we’re going to do that next time.

What Are We Doing?

In the book, I’ve rewritten the problem descriptions for each problem in my own style, which hopefully makes them all a bit more consistent and streamlined for readers. I think it’s also important to read original problem descriptions too, though, so you can practice getting used to different people’s ways of describing problems and extracting what’s most important.

So, if you haven’t read the problem description above, please do that now :) try to understand what the problem is asking, and then compare with what I’m about to say next:

What’s going on here is that we have $N$ cows numbered $1, 2, \ldots, N$. We wanted to arrange them in a specific order (that’s the $A$ list), but that $A$ list was stolen, so we don’t know that order anymore.

What we do have is the $B$ list, which tells us the sums of consecutive values in the stolen $A$ list. For example, $B[1]$ tells us the sum of $A[1]$ and $A[2]$, $B[2] = A[2] + A[3]$, and so on. So any value $B[i]$ tells us what $A[i] + A[i + 1]$ is.

Be careful indexing these lists: the problem talks about $A$ and $B$ as starting with index 1, not index 0. So in our code to follow, to match that 1-based indexing, we’re going to start using our lists at index 1 as well (we’ll just put a placeholder value to waste index 0).

Notice that $B$ has one fewer value than $A$. That’s because, for example, if $A$ had three values, then we only need two $B$ values to get us to the end of $A$: $B[1] = A[1] + A[2]$ and $B[2] = A[2] + A[3]$.

What we need to do is come up with an $A$ list that could have produced the $B$ list. For the $A$ list to be valid, it has to have exactly $N$ values, and each value from $1$ to $N$ has to be used exactly once (that’s the definition of a permutation).

An Example

Let’s run through an example of what our program is supposed to do.

Here’s some sample input:

6
9 5 7 9 5

This tells us that the $A$ list that we want to recover has 6 values in it (and that those values must be a permutation of the numbers from 1 to 6). And the $B$ list is $[9, 5, 7, 9, 5]$.

Now, I’m going to give you an $A$ list that I claim is the answer. Here it is…

Answer: $[5, 4, 1, 6, 3, 2]$

First of all, this really is a permutation of the numbers from 1 to 6. All six of those numbers are in there exactly once. And second, we can confirm that the $B$ list that we were given really could have been produced from this $A$ list. If that’s clear to you already then skip to the next heading, but otherwise:

Now, what I find really interesting here is that there’s actually another possible answer… and here it is!

Another answer?: $[6, 3, 2, 5, 4, 1]$

Try it, and you should see that this list is compatible with $B$ as well. So what do we do? Which answer is correct? In a situation like this where we can have multiple answers for $A$, the problem description tells us to pick the minimum permutation, which just means the one with the smaller initial value. That’s why we pick the one that starts with the $5$, not the one that starts with the $6$.

So for the record, here’s the correct output:

5 4 1 6 3 2

OK, now let’s pull in the stuff from Chapter 9 (our Complete Search chapter). For those who haven’t read it in a while or don’t want to read it :) I’m going to summarize what we need here.

A Complete Search algorithm involves going through all candidate solutions and keeping track of the best one that we find. That “best one” is our answer.

One of our examples in Chapter 9 was a problem about lifeguards: we hired $N$ lifeguards, each who monitored a pool for their shift – but, uh oh, budget cuts, we have to fire one! We want to keep pool coverage as high as possible, so we want to fire the lifeguard that has minimal effect on our coverage. And how did we do it in Chapter 9? We tried firing each lifeguard separately, then calculating the pool coverage for each choice. Once we try them all, we can choose the best one to fire. (In a secret bit of information, I can promise you that the fired lifeguard is OK – it turns out that they preferred being in a rock band anyway, so they started a band named Discharged Lifeguards and made it big. :D )

Now, why might we think that Complete Search could work for our Photoshoot problem?

Well, check this out: suppose that someone saw you working on the Photoshoot problem, and as they walked past you, they whispered “psssst, the $A$ list starts with 5”. Then they’re gone. (They’re off to their new lifeguard job.)

What’s cool is that only knowing the starting point of $A$, you can reconstruct the rest of $A$ perfectly, with no mistakes. You don’t need any more information from the whispering person.

For example, $B[1]$ tells us that $A[1] + A[2]$ is 9. But our whispering person already told us that $A[1]$ is 5. So then $A[2]$ must be 4! There’s nothing else that $A[2]$ could be if we want to get a total of 9. We’re just subtracting here.

And now that we know $A[2]$, we can figure out $A[3]$ perfectly as well. We know from $B[2]$ that $A[2] + A[3]$ is 5. We already know that $A[2]$ is 4, so $A[3]$ must be $5 - 4 = 1$. The subtraction again gives us the remainder for what $A[3]$ must be. We’re going to use this subtraction in our code later.

And we can keep going until we get to the end of $A$.

Now, of course, there’s no whispering person telling us the first value of $A$. But that’s OK, because we have Complete Search! Yep, all we do is try each possible value for $A[1]$.

First we try setting $A[1]$ to 1, build out the whole array and ask, “is this right?”. In this case it won’t be… starting with $A[1] = 1$, we eventually end up with the $A$ list $[1, 8, -3, 10, -1, 6]$ which isn’t a permutation of the numbers 1-6 at all. So we gotta continue.

So next we try setting $A[1]$ to 2 and playing that out. We’ll end up with a mess again. We then set $A[1]$ to 3, then 4, then 5… oh, there we go, 5 works! And since we want the minimum permutation, this is our answer. We stop here and output what we’ve found.

Do We Have a Valid Permutation?

As we systematically try starting our proposed $A$ list with 1, then 2, then 3, and so on, we’re going to generate a lot of $A$ lists that don’t work; i.e. $A$ lists that are not valid permutations. It’s very helpful with complete search algorithms to write a function to tell us something about our current attempted solution. (In the Lifeguards case, it was: what’s the coverage of the pool with this specific lifeguard fired?) Here, we need to know if our current $A$ list is a valid permutation or not.

To determine if it’s a valid permutation, we’re going to use a set to keep track of which values we’ve seen so far. We never want to see the same value twice – if we do, then it isn’t a valid permutation. We also need to check that we don’t have any illegal numbers in our list; i.e. we don’t want numbers less than 1 or greater than the number of elements in the list (those can’t be part of the permutation).

Here’s our code for this helper function.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

# Convention in this program:
# In a list, index 0 is ignored (set to 0 placeholder value).
# The A and B lists in the problem description start
# at A[1] and B[1], so I wanted our lists to start at 1 as well.


def is_valid_permutation(lst):
    """
    lst is a list of integers.
    The integers are at lst[1], lst[2], ..., lst[n].

    Return True if lst is a permutation of the integers from 1 to n, False otherwise.
    """
    already_used = set()
    for i in range(1, len(lst)):
        # Already used? Out of bounds? Both bad!
        if lst[i] in already_used or lst[i] < 1 or lst[i] >= len(lst):
            return False
        already_used.add(lst[i])
    return True

With that helper function, we’re in good shape to implement our complete search.

Here’s the code for our main program. I’m using possible_a_lst and b_lst instead of $A$ and $B$, for better Python variable names.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

# This program uses files; see Chapter 7 for review
input_file = open('photo.in')
output_file = open('photo.out', 'w')

n = int(input_file.readline())
b_lst = input_file.readline().split()
for i in range(len(b_lst)):
    b_lst[i] = int(b_lst[i])
b_lst = [0] + b_lst  # Per convention, put placeholder at index 0

# Complete search over all possible values for possible_a_lst[1]

start_val = 1
while True:
    possible_a_lst = [0, start_val]  # Placeholder at index 0 again
    for i in range(2, len(b_lst) + 1):  # Generate rest of A list
        next_value = b_lst[i - 1] - possible_a_lst[i - 1]
        possible_a_lst.append(next_value)
    if is_valid_permutation(possible_a_lst):
        break  # Found answer!
    start_val = start_val + 1  # Continue the complete search

output_lst = []
for val in possible_a_lst[1:]:
    output_lst.append(str(val))

output_file.write(' '.join(output_lst) + '\n')

input_file.close()
output_file.close()

That start_val variable (13) keeps track of the solution that we’re trying now. Setting start_val to 1, for example, means that we’re about to try a solution for $A$ where $A[1] = 1$. Setting start_val to 2 means that we’re about to try $A[1] = 2$, and so on. The while loop (14) continues until we successfully find a permutation.

Line (17) is where we figure out the next value in the $A$ list. To do that, we subtract the previous $A$ value from the previous $B$ value – that’s the subtraction we talked about earlier.

Coming Up Next

Today’s solution was fast enough for the time limit – try it! :)

If you’ve read Chapter 10 of the book, all about efficiency of algorithms, then you might recognize that our program is an $O(n^2)$ algorithm. Our is_valid_permutation function takes $O(n)$ time, and we might call that thing a total of $O(n)$ times as well, for a total of $O(n^2)$ work.

At the start of this post, I promised that there was a way to solve this in only $O(n)$ time. This means that we cannot afford to try $n$ possible candidates for $A[1]$ – that’s too slow! We have to be able to dramatically cut down on the number of possible permutations we need to try. And the way to do that is …

Next time :)

More Practice

There’s more! Please check the list of exercises that I’m curating for the book for even more practice with the material from the book.

Happy learning!

Written on February 11th, 2025 by Daniel Zingaro
Feel free to share!

You may also enjoy: