Longest spell to cast from pages of spellbook follow-up

Question

This question is from the PCTC 2022 R2 Past Paper and is a follow-up on my previous question. Previous question

I have implemented several solutions suggested, such as creating an array with pages already visited (@J_H) and ignoring dead-ends (self-referring pages) (@aghast), as well as stopping the loop if the algorithm is already re-tracing the same spell (@leo). However, it still exceeds the time limit. Is there anything wrong with my approach/more efficient ways for loop detection?

---

Problem

Finn is a wizard’s apprentice. While clearing up his master’s attic, he found an old spellbook and he decided to use it himself.

The spellbook has n pages, labelled 1 to n. On each page there is some arcane language, followed by a number.

To cast a spell, Finn needs to open the spellbook on a page of his choice, utter the words on the page and then turn the book to the page whose number is given there (which could be itself). He then utters the words on that page and follows the number, and so on…

He can use each page only once, so if he is directed to a page already read, he must end the spell there.

Finn would like to cast the longest spell available in this book.

As trial-and-error is not a safe way of casting spells, he has asked you to write a program to find the length of the longest spell he could cast and also how many of this length there are.

Input Format

The first line contains a positive whole number n – the number of pages in the spellbook. Each of the following n lines contains a single integer between 1 and n inclusive. The ith of these numbers is the page number which page #i redirects to.

Output Format

On the first line, output the total number of pages included in the longest spell. On the second line, output the number of possible spells of this length.

Example Input

6
3
5
5
3
1
6

Example Output

4
2

Example Explanation

In this book of 6 pages (the first input), there are two possible spells of length 4.

One of them starts on page 2 and goes 2 -> 5 -> 1 -> 3, and the next one starts on page 4 and goes 4 -> 3 -> 5 -> 1.

No other starting page produces a longer spell. Starting on page 1, 3 or 5, it is only possible to cast a spell of length 3 before returning to the starting page. Starting on page 6 immediately returns to itself.

Hence the output is 4 (the longest spell), followed by 2 (the number of such spells).

Constraints

n will satisfy 2 <= n <= 300,000 note: programs must complete in under 3 seconds and use less than 64Mb of memory note: solutions which run out of memory are usually reported as 'wrong answer'; solutions that run out of time are reported as 'time limit exceeded'

My Code

n = int(input())
og = [0]*n
visited_pages = [0]
total_visited_pages = [-1]*n #-1 is a marker for an empty element
spell_lengths = [0]
for current_page_number in range(n):
    og[current_page_number] = int(input()) #copying input into list
for current_page_number in range(len(og)):
    loop_detected = False
    j = current_page_number
    visited_pages.clear()
    if j == og[j]-1: #ignores any dead ends e.g page 1 has 1 on it
        spell_lengths.append(1)
    else:
        while j not in visited_pages and loop_detected == False:
            if j in total_visited_pages: #check if this loop has been traced before
                for x in range(n):
                    if j in total_visited_pages[x]:
                        loop_detected = True #breaks the while loop as a loop has been detected
                        spell_lengths.append(len(total_visited_pages[x])) #append the length since all loops have the same length
            visited_pages.append(j) #we have visited j
            j = og[j]-1 #since 0 is the first element of an array, -1 from the number
        spell_lengths.append(len(visited_pages)) #add the spell length to the list
        total_visited_pages[current_page_number] = visited_pages #add to the total visited pages
print(max(spell_lengths))
print(spell_lengths.count(max(spell_lengths)))

Results on Website

Today 12:01:27, Results

Testset Outcome Score   Testcases
t0      Passed  0 / 0   Testcase 1: Ok
t2      Passed  1 / 1   Testcase 1: Ok
t3      Passed  1 / 1   Testcase 1: Ok
t4      Passed  1 / 1   Testcase 1: Ok
t5      Passed  1 / 1   Testcase 1: Ok
t6      Error   0 / 1   Testcase 1: Time limit exceeded
t7      Error   0 / 1   Testcase 1: Time limit exceeded
t8      Error   0 / 1   Testcase 1: Time limit exceeded
t9      Error   0 / 1   Testcase 1: Time limit exceeded
t10     Error   0 / 1   Testcase 1: Time limit exceeded
t1      Passed  1 / 1   Testcase 1: Ok

Answer 1

Naming

Most of your names are descriptive enough, but I can't, for the life of me, figure out what og stand for… pages or book should be more appropriate. In a lesser way, j might benefit from a better name as well.

Offset

In two places, you need to subtract 1 from the page number that you need to read next in order to compensate between the 0-based list indexes and the 1-based page numbers. You most likely want to do this only once, preferably while parsing the input, in order to reduce the cognitive burden necessary to understand that logic.

Not quite the behavior that you expect

As it currently stand, your total_visited_pages does nothing to help speed up the program. It is a waste of space at best, a waste of time at worst. Let's print its value at the end of each loop for the provided test case:

[[0, 2, 4], -1, -1, -1, -1, -1]
[[1, 4, 0, 2], [1, 4, 0, 2], -1, -1, -1, -1]
[[2, 4, 0], [2, 4, 0], [2, 4, 0], -1, -1, -1]
[[3, 2, 4, 0], [3, 2, 4, 0], [3, 2, 4, 0], [3, 2, 4, 0], -1, -1]
[[4, 0, 2], [4, 0, 2], [4, 0, 2], [4, 0, 2], [4, 0, 2], -1]
[[], [], [], [], [], -1]

The culprit being both visited_pages.clear() and that an affectation (total_visited_pages[current_page_number] = visited_pages) will never copy the data but only share references. As it stand, you're storing multiple references to the same visited_pages into total_visited_pages which makes it completely useless since you already know that j is not in there.

The fix is simple as you need to change the clear call into creating a completely new list with visited_pages = [] and now you get the caching behavior that you were expecting:

[[0, 2, 4], -1, -1, -1, -1, -1]
[[0, 2, 4], [1, 4, 0, 2], -1, -1, -1, -1]
[[0, 2, 4], [1, 4, 0, 2], [2, 4, 0], -1, -1, -1]
[[0, 2, 4], [1, 4, 0, 2], [2, 4, 0], [3, 2, 4, 0], -1, -1]
[[0, 2, 4], [1, 4, 0, 2], [2, 4, 0], [3, 2, 4, 0], [4, 0, 2], -1]
[[0, 2, 4], [1, 4, 0, 2], [2, 4, 0], [3, 2, 4, 0], [4, 0, 2], -1]

I'd also posit that even with this fix, the total_visited_pages is still useless as it contains either lists or -1s, so the check if j in total_visited_pages will always be false anyways. More on how to fix that later.

Tracking the longest spell

Interestingly enough, using total_visited_pages you already have all the information you need to re-create the spell_lengths list: you just need to take the len of each list in total_visited_pages (and account for the self-referencing pages) and you're done. But you can do better as you don't even need to store all this information. You only care about the maximum length and the amount of spells of said maximum length.

So you can create a small helper to avoid searching for the max after the facts:

from dataclasses import dataclass


@dataclass()
class MaxTracker:
    maximum: int = 0
    amount: int = 0

    def track_value(self, value):
        if value > self.maximum:
            self.maximum = value
            self.amount = 1
        elif value == self.maximum:
            self.amount += 1

So instead of appending to a list, you just track_value on your MaxTracker and you can easily retrieve the information you’re interested in using simple attributes lookup.

Using functions

In order to more easily test and debug your code, you should put it into functions. It also helps when reasoning about the logic by breaking it down into meaningful components. It is also helpful when dealing with simple benchmarks.

Applying the previous items to your code and throwing better iteration habits into the mix could yield:

from dataclasses import dataclass, astuple


@dataclass()
class MaxTracker:
    maximum: int = 0
    amount: int = 0

    def track_value(self, value):
        if value > self.maximum:
            self.maximum = value
            self.amount = 1
        elif value == self.maximum:
            self.amount += 1


def parse_book(book):
    total_visited_pages = [-1 for _ in book]  # -1 is a marker for an empty element
    spell_lengths = MaxTracker()
    for page_number, link_to in enumerate(book):
        if page_number == link_to:
            spell_lengths.track_value(1)
        else:
            loop_detected = False
            current_page_number = page_number
            visited_pages = []
            while current_page_number not in visited_pages and not loop_detected:
                if current_page_number in total_visited_pages:  # check if this loop has been traced before
                    for already_visited in total_visited_pages:
                        if current_page_number in already_visited:
                            loop_detected = True  # breaks the while loop as a loop has been detected
                            spell_lengths.track_value(len(already_visited))  # append the length since all loops have the same length
                visited_pages.append(current_page_number)  # we have visited the current page
                current_page_number = book[current_page_number]
            spell_lengths.track_value(len(visited_pages))  # check if the spell length is the longest one yet
            total_visited_pages[page_number] = visited_pages  # add to the total visited pages

    return spell_lengths


def main():
    page_count = int(input())
    # Convert 1-based pages numbers into 0-based list index
    print(*astuple(parse_book([int(input()) - 1 for _ in range(page_count)])), sep='\n')


if __name__ == '__main__':
    main()

This will allow you to do things like (assuming the code is in spells.py):

$ python
Python 3.11.8 (main, Feb 12 2024, 14:50:05) [GCC 13.2.1 20230801] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import timeit
>>> timeit.timeit('test(book)', 'from spells import parse_book as test; pages=[3, 5, 5, 3, 1, 6]; book=[p - 1 for p in pages]')
3.349241663000612

Performances

The instructions indicate that your book can have up to 300k pages, so let's increase the size of our test cases accordingly and try:

timeit.timeit('parse_book(book)', 'from spells import parse_book; book=list(range(1, 100000)) + [0]')

No, this is taking too long, let's do only a single pass into the parse_book function:

timeit.timeit('parse_book(book)', 'from spells import parse_book; book=list(range(1, 100000)) + [0]', number=1)

Still taking ages, maybe if we reduced the amount of pages:

timeit.timeit('parse_book(book)', 'from spells import parse_book; book=list(range(1, 10000)) + [0]', number=1)

Nope, more than ½h and still nothing…

timeit.timeit('parse_book(book)', 'from spells import parse_book; book=list(range(1, 1000)) + [0]', number=1)

At last, a mere 12.85 seconds… For a program supposed to run in less than 3 seconds… using a hundredfold input… There must be a better way.

While not providing you with such a better way, I’ll just point out that your use of lists makes your program cripplingly slow as each in test will perform a linear scan of the list. And you get to perform it for each page, and each link (j) inside that page, which make your overall performance \$O(n^3)\$. This is bad. Instead, you could use a set for your visited_pages so existence checks will be \$O(1)\$ instead of \$O(n)\$ and a dictionary for your total_visited_pages so retrieving a cycle from an already processed page is \$O(1)\$ as well:

from dataclasses import dataclass, astuple


@dataclass()
class MaxTracker:
    maximum: int = 0
    amount: int = 0

    def track_value(self, value):
        if value > self.maximum:
            self.maximum = value
            self.amount = 1
        elif value == self.maximum:
            self.amount += 1


def parse_book(pages):
    spell_lengths = MaxTracker()
    total_visited_pages = {}
    for page_num, link_to in enumerate(pages):
        if page_num == link_to:
            spell_lengths.track_value(1)
        else:
            visited_pages = set()
            current_page = page_num
            while current_page not in visited_pages:
                if current_page in total_visited_pages:
                    visited_pages.update(total_visited_pages[current_page])
                    break
                visited_pages.add(current_page)
                current_page = pages[current_page]
            spell_lengths.track_value(len(visited_pages))
            total_visited_pages[page_num] = visited_pages

    return spell_lengths


def main():
    page_count = int(input())
    # Convert 1-based pages numbers into 0-based list index
    pages = [int(input()) - 1 for _ in range(page_count)]
    print(*astuple(parse_book(pages)), sep='\n')


if __name__ == '__main__':
    main()

This allows for a faster execution time, but we’re still far from the expected one, tough:

$ python
Python 3.11.8 (main, Feb 12 2024, 14:50:05) [GCC 13.2.1 20230801] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import timeit
>>> timeit.timeit('parse_book(book)', 'from spells import parse_book; book=list(range(1, 10000)) + [0]', number=1)
6.2739594989998295
>>> timeit.timeit('parse_book(book)', 'from spells import parse_book; book=list(range(1, 100000)) + [0]', number=1)

…aaaand that last test completely filled up my 16G of RAM…

There must be a better way!

I didn't want to write this section, at first. But for the sake of the FutureReader™, let's dive into how to overcome these last few obstacle. I recommend you pause first and try to understand the underlying issues with the previous code before reading the full answer.

There are two major problems with the previous solution:

1. We store the loop in the total_visited_pages table only for the starting page at once, so we need to "re-discover" the loop for each page that is part of it;
2. We store the whole loops as the values of the total_visited_pages table, wasting space and time to construct them, but we ultimately only care about their length.

A last issue, that is more of a bug, is that, regardless of where the loop start in the spell, we store the whole sequence into total_visited_pages and treat it as a loop afterwards. It is necessary to take into account the fact that not all loops get back to the first page of the spell when computing the length of both the loops and the new spells leading to this loop.

Speaking of spells leading to loops, we need to understand where the loops are created and where they are merely reached from a new starting page. To this extend, we must recognize that the test if current_page in total_visited_pages in the loop is merely reaching an existing loop from a new starting page; and the natural (current_page in visited_pages) way of exiting the loop is detecting a new loop in the current spell. Luckily Python lets us easily differentiate between these two exits of the loop by using the while…else construct. So it is now easy to only add the linear length of the new spell before breaking out of the loop, and handle the spells loops in the else clause of the while loop.

But we must take a step back first and remember that not all pages of a spell are part of a loop and check which page is the first of the loop in order to apply this separation (linear path + loop) when adding the new loop to the total_visited_pages tables as well. To simplify things and still keep performant, \$O(1)\$ existence checks of sets, we need to remember that they also remember their insertion order (sets are merely dicts with no values) and we can iterate on them as though we have a list (or convert them to list if it is more convenient for methods such as index).

Taking all that into consideration lead to the following code:

from dataclasses import dataclass, astuple


@dataclass()
class MaxTracker:
    maximum: int = 0
    amount: int = 0

    def track_value(self, value):
        if value > self.maximum:
            self.maximum = value
            self.amount = 1
        elif value == self.maximum:
            self.amount += 1


def parse_book(pages):
    spell_lengths = MaxTracker()
    total_visited_pages = {}
    for page_num, link_to in enumerate(pages):
        if page_num == link_to:
            total_visited_pages[page_num] = 1
        elif page_num not in total_visited_pages:
            visited_pages = set()
            current_page = page_num
            while current_page not in visited_pages:
                if current_page in total_visited_pages:
                    # We reached an existing spell loop
                    spell_length = len(visited_pages) + total_visited_pages[current_page]
                    spell_lengths.track_value(spell_length)
                    for index, page_number in enumerate(visited_pages):
                        total_visited_pages[page_number] = spell_length - index
                    break
                # Turn the book to the right page and remember where we went from
                visited_pages.add(current_page)
                current_page = pages[current_page]
            else:
                # We found a new spell loop
                spell_length = len(visited_pages)
                spell_lengths.track_value(spell_length)
                visited_pages = iter(visited_pages)
                # Process the linear part before the loop first
                for index, page_number in enumerate(visited_pages):
                    if page_number == current_page:
                        loop_index = index
                        break
                    total_visited_pages[page_number] = spell_length
                    spell_length -= 1
                # Handle the loop next and remember that each component of the loop has the same length
                for page_number in visited_pages:
                    if not loop_index:
                        spell_lengths.track_value(spell_length)
                    total_visited_pages[page_number] = spell_length

    return spell_lengths


def main():
    page_count = int(input())
    # Convert 1-based pages numbers into 0-based list index
    pages = [int(input()) - 1 for _ in range(page_count)]
    print(*astuple(parse_book(pages)), sep='\n')


if __name__ == '__main__':
    main()

And this version is not only space efficient, it is also pretty fast compared to the previous ones:

$ python
Python 3.11.8 (main, Feb 12 2024, 14:50:05) [GCC 13.2.1 20230801] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import timeit
>>> timeit.timeit('parse_book(pages)', 'from spells import parse_book; pages = list(range(1_000_000))', number=1)
0.08444321500019214
>>> timeit.timeit('parse_book(pages)', 'from spells import parse_book; pages = list(range(1, 1_000_000)) + [0]', number=1)
0.26171460300020044
>>> timeit.timeit('parse_book(pages)', 'from spells import parse_book; pages = [1_000_000] + list(range(1_000_000))', number=1)
0.26340376899997864