Problem Statement in English

You’re given an array of positive integers nums.

You need to find the maximum size of a subset that can be rearranged into the following pattern:

$$ [x, x^2, x^4, \dots, x^k, \dots, x^4, x^2, x] $$

In other words, the numbers keep getting squared until they reach a single peak, and then appear in reverse order.

Return the maximum possible length of such a subset.

Approach

The important observation is that every number except the peak must appear at least twice: once on the left side of the pattern and once on the right. The peak only needs to appear once.

So, for each possible starting number, we repeatedly square it and extend the two sides of the pattern for as long as the current number occurs at least twice.

There are a few details we need to handle.

1. Handle 1 Separately

The number 1 is a special case because $1^2 = 1$.

A valid subset made entirely of ones must have an odd length. Therefore, we take the largest odd number that does not exceed the number of ones:

res = (freq.pop(1, 0) - 1) | 1

The expression (count - 1) | 1 rounds the count down to the nearest odd number. For example, if there are 4 ones, we can use 3 of them. If there are no ones, res starts at -1 and will be replaced when we process the other numbers.

We also remove 1 from the frequency map because repeatedly squaring it would never move us to a new value.

2. Skip Chains That Start in the Middle

Suppose both 2 and 4 occur at least twice. There is no reason to start a chain from 4, because any chain beginning with 4 is shorter than the chain beginning with 2:

4 -> 16
2 -> 4 -> 16

For each number x, we check whether it is a perfect square and whether its square root occurs at least twice. If both conditions are true, x is already part of a chain with an earlier starting point, so we skip it.

sq = isqrt(x)
if sq * sq == x and freq.get(sq, 0) > 1:
    continue

3. Build Both Sides of the Pattern

We now keep squaring x while there are at least two copies of it:

n = 0
while x < 31623 and freq.get(x, 0) > 1:
    n += 2
    x *= x

Each level adds 2 to the length because one copy goes on either side of the peak.

The values in nums are at most $10^9$. Since $31623^2 > 10^9$, squaring a number greater than or equal to 31623 cannot produce another value in the input.

4. Choose the Peak

When the loop stops, there are two possibilities:

  • x exists once. We can use it as the peak, so the chain length is n + 1.
  • x does not exist. The previous value must be the peak. We counted that value twice during the last iteration, so the chain length is n - 1.

The code handles both cases with:

n + ((x in freq) << 1) - 1

x in freq is either True, which behaves like 1, or False, which behaves like 0. The expression therefore adds 1 when x exists and subtracts 1 when it doesn’t.

Solution in Python


class Solution:
    def maximumLength(self, nums: list[int]) -> int:
        freq = Counter(nums)
        res = (freq.pop(1, 0) - 1) | 1

        for f in freq:
            x = f
            sq = isqrt(x)
            if sq * sq == x and freq.get(sq, 0) > 1:
                continue

            n = 0
            while x < 31623 and freq.get(x, 0) > 1:
                n += 2
                x *= x

            res = max(res, n + ((x in freq) << 1) - 1)

        return res

Complexity

  • Time: $O(n + u \log \log M)$
    where $n$ is the length of nums, $u$ is the number of unique values, and $M$ is the maximum value in nums. Building the frequency map takes $O(n)$.

    The $\log \log M$ comes from how quickly x grows inside the loop. If the first value is $x_0$, repeatedly squaring it gives us:

    • Iteration 0: $x_0$
    • Iteration 1: $x_0^2$
    • Iteration 2: $(x_0^2)^2 = x_0^4 = x_0^{2^2}$
    • Iteration 3: $(x_0^4)^2 = x_0^8 = x_0^{2^3}$

    Therefore, after $k$ iterations, the current value is $x_0^{2^k}$. The loop can only continue while this value is below $M$, so we find the point where it reaches the limit:

    $$ x_0^{2^k} \ge M $$

    We need to isolate $k$, which is currently inside an exponent inside another exponent. Taking $\log_2$ once brings the outer exponent down:

    $$ 2^k \log_2 x_0 \ge \log_2 M $$

    Dividing by $\log_2 x_0$ gives us:

    $$ 2^k \ge \frac{\log_2 M}{\log_2 x_0} $$

    We then take $\log_2$ a second time to bring down the remaining exponent and isolate $k$:

    $$ k \ge \log_2 \left(\frac{\log_2 M}{\log_2 x_0}\right) $$

    To find the maximum possible number of iterations, we use the smallest possible starting value. A smaller $x_0$ grows more slowly and therefore takes more squarings to reach $M$. Since 1 is handled separately, the smallest possible value here is $x_0 = 2$. Plugging it in gives us:

    $$ k \ge \log_2 \left(\frac{\log_2 M}{\log_2 2}\right) = \log_2 \log_2 M $$

    Therefore, a chain can contain at most $O(\log \log M)$ values.

  • Space: $O(u)$
    for the frequency map.

And we are done.