LeetCode 481: Magical String

Problem Restatement

We are given an integer n.

There is a special infinite string s made only of characters '1' and '2'.

The first part of the string is:

1221121221221121122...

This string is called magical because if we group equal consecutive characters, the group lengths form the same string again.

For example:

s      = 1 22 11 2 1 22 1 22 11 2 11 22 ...
length = 1 2  2  1 1 2  1 2  2  1 2  2  ...

The length sequence is again:

122112122122...

We need to return how many '1' characters appear in the first n characters of s.

The constraints are 1 <= n <= 10^5. The examples include n = 6, where the first six characters are "122112" and contain three '1' characters.

Input and Output

Function shape:

def magicalString(n: int) -> int:
    ...

Examples

Example 1:

n = 6

The first six characters are:

122112

There are three '1' characters.

Answer:

3

Example 2:

n = 1

The first character is:

1

Answer:

1

First Thought: Build the String Directly

The string describes its own run lengths.

We already know the beginning:

122

This means:

Now the next length comes from the next unread character in the string.

Starting with:

s = 122
      ^
      read pointer

The read pointer is at index 2, where s[2] = 2.

The previous group used character 2, so the next group must use character 1.

Since the count is 2, append two 1s:

s = 12211

Move the read pointer forward.

Now the next count is s[3] = 1.

The next character alternates to 2.

Append one 2:

s = 122112

This process continues until the string has at least n characters.

Key Insight

Use the current contents of s as instructions for how to extend s.

Maintain:

At every step:

count = s[i]

Then append next_num exactly count times.

After that, flip next_num:

next_num = 3 - next_num

This works because the only possible values are 1 and 2.

Algorithm

Start with:

s = [1, 2, 2]
i = 2
next_num = 1

Then repeat while len(s) < n:

1. Read the next count from s[i].
2. Append next_num that many times.
3. Flip next_num.
4. Move i forward.

After the string reaches length n, count how many 1s appear in s[:n].

Correctness

The magical string is defined by its run lengths.

The initial string [1, 2, 2] gives the first two groups:

1
22

The pointer i = 2 points to the length of the third group.

At each step, s[i] is the length of the next group to append. Since groups alternate between 1 and 2, next_num is always the correct group value.

The algorithm appends exactly s[i] copies of next_num, so the new group has the required length.

Then it advances to the next run-length instruction and flips the appended digit for the following group.

By induction, after every iteration, the generated prefix agrees with the magical string. Once the generated prefix has length at least n, the first n characters are correct. Counting the 1s in that prefix gives the required answer.

Complexity

The loop may append one or two values each step, but the total number of appended values is proportional to n.

Implementation

class Solution:
    def magicalString(self, n: int) -> int:
        s = [1, 2, 2]

        i = 2
        next_num = 1

        while len(s) < n:
            count = s[i]

            for _ in range(count):
                s.append(next_num)

            next_num = 3 - next_num
            i += 1

        return s[:n].count(1)

Code Explanation

We seed the known prefix:

s = [1, 2, 2]

The read pointer starts at index 2:

i = 2

That value tells us the length of the third group.

The third group should be made of 1s:

next_num = 1

Inside the loop, this reads the next group length:

count = s[i]

Then we append the next digit count times:

for _ in range(count):
    s.append(next_num)

After creating a group of 1s, the next group must be 2s. After creating a group of 2s, the next group must be 1s.

This flip is done by:

next_num = 3 - next_num

Finally, we move to the next instruction:

i += 1

Return the number of 1s in the first n generated characters:

return s[:n].count(1)

Slightly Optimized Implementation

We can count 1s while building the string.

This avoids scanning the prefix again at the end.

class Solution:
    def magicalString(self, n: int) -> int:
        if n <= 3:
            return 1

        s = [1, 2, 2]
        ones = 1

        i = 2
        next_num = 1

        while len(s) < n:
            count = s[i]

            for _ in range(count):
                if len(s) == n:
                    break

                s.append(next_num)

                if next_num == 1:
                    ones += 1

            next_num = 3 - next_num
            i += 1

        return ones

Testing

def run_tests():
    s = Solution()

    assert s.magicalString(1) == 1
    assert s.magicalString(2) == 1
    assert s.magicalString(3) == 1
    assert s.magicalString(4) == 2
    assert s.magicalString(5) == 3
    assert s.magicalString(6) == 3
    assert s.magicalString(10) == 5

    print("all tests passed")

run_tests()

Test meaning: