ZigZag Conversion en LeetCode: Solución optimizada en Python


Continuando con los problemas de LeetCode llegué al problema de la conversión ZigZag que parece ser un poco más sencillo que el anterior de los palíndromos. Este problema dice si tenemos el texto "PAYPALISHIRING" distribuido en zigzag en una cantidad dada de renglones, por ejemplo 3, el zigzag sería:

P A H N
APLSIIG
Y I R  

Entonces al concatenar línea por línea la salida sería "PAHNAPLSIIGYIR".

Al analizar un poco el problema mi primera idea fue simular la matriz, así que tomé el string principal y construí subcadenas que representan las columnas de la matriz:

P A H N
APLSIIG
Y I R  

Con esta idea en mente escribí:


class Solution:
    def convert(self, s: str, numRows: int) -> str:
        if 2 == numRows:
            return s[::2] + s[1::2]

        chunks = []
        i = 0
        direction = 0 # 0 = down, 1 = up
        while i < len(s):
            if 0 == direction:
                to = i + numRows
                chunks.append(s[i:to])
                i = to
                direction = 1
                fillerSize = numRows - 1
            else:
                fillerSize -= 1
                to = i + 1
                zagFiller = ' ' * fillerSize
                chunks.append(zagFiller + s[i:to])
                i = to

                if 1 >= fillerSize:
                    direction = 0

        chunkLen = len(chunks[0])
        result = ''
        for i in range(chunkLen):
            result += ''.join([chunk[i] for chunk in chunks if len(chunk) > i])
        return result.replace(" ", "")

Tuve problemas para completar esta solución porque los escenarios de numRows = 2 no se resolvían en esta lógica, por lo que tuve que agregar la condición if 2 == numRows:. Pero una vez completo construía lo siguiente: ['PAY', ' P ', 'ALI', ' S ', 'HIR', ' I ', 'NG'] para luego concatenar las primeras posiciones de cada subcadena, luego la segunda posición de cada subcadena, etc. Le pasé mi solución a Claude y me dijo lo siguiente:

Construyes chunks que simulan la columna vertical y las diagonales del zigzag, los alineas con padding de espacios, y luego lees por filas. Funciona, pero estás creando strings intermedios con padding (' ' * fillerSize + s[i:to]), lo que significa que el espacio que usas no es solo O(n) para el resultado, sino potencialmente más por todo el padding que construyes y luego descartas con .replace(" ", ""). En el peor caso con muchas filas, ese overhead puede ser significativo. Además el replace(" ") es un hack que puede romper el output si el string original contiene espacios. El problema dice que s solo tiene letras y dígitos, así que no falla en los tests, pero es frágil por diseño.

A partir de los comentarios cambié de enfoque: en lugar de simular las columnas ahora simularía filas:


class Solution:
    def convert(self, s: str, numRows: int) -> str:
        if len(s) == 1:
            return s
        center = -1
        half = numRows // 2 - 1
        result = ''
        diag = numRows - 2
        zag = diag if diag > 0 else 0

        if 0 != numRows % 2:
            center = numRows // 2

        for i in range(numRows):
            if i == 0 or i == numRows - 1:
                result += s[i::numRows + zag]
            elif i == center:
                result += s[i::numRows - 1]
            else:
                fs = (numRows - i - 1) * 2
                ss = i * 2
                result += self.side(s, i, fs, ss)
        return result
    
    def side(self, s, i, fs, ss):
        result = ''
        skips = [fs, ss]
        skipPos = False
        while i < len(s):
            result += s[i]
            i += skips[int(skipPos)]
            skipPos = not skipPos
        return result

La idea de este enfoque es calcular la posición de cada letra de la cadena original en la cadena resultante. Esta solución se murió cuando la string original solo tiene un caracter y por eso agregué la condición if len(s) == 1:. Sobre este enfoque Claude me dijo:

Buen progreso — ya estás pensando en términos de filas y de saltos entre caracteres, que es el insight correcto. Lo que te pregunto para que reflexiones:
  1. Tu manejo de center para filas impares — ¿por qué necesitas ese caso especial? ¿No debería estar cubierto por la lógica general de side()?
  2. Mira tus valores de fs y ss para una fila intermedia cualquiera — ¿cuál es la relación entre esos dos valores? ¿Suman siempre lo mismo? Si es así, ¿qué representan geométricamente en el zigzag?

Con esas reflexiones acompañadas con una pequeña sesión de preguntas aclaratorias logré eliminar los casos especiales en el for, quedando:


class Solution:
    def convert(self, s: str, numRows: int) -> str:
        if len(s) == 1 or numRows == 1:
            return s
        result = ''

        for i in range(numRows):
            fs = (numRows - i - 1) * 2
            ss = i * 2
            result += self.side(s, i, fs, ss)
        return result
    
    def side(self, s, i, fs, ss):
        result = ''
        fs = fs if fs > 0 else ss
        ss = ss if ss > 0 else fs
        skips = [fs, ss]
        skipPos = False
        while i < len(s):
            result += s[i]
            i += skips[int(skipPos)]
            skipPos = not skipPos
        return result

Los comentarios de Claude sobre esta optimización son:

Eliminaste el caso especial del for moviendo la lógica a side(). ¿Cambió algo respecto a tu versión anterior? No asintóticamente — ambas son O(n) tiempo y O(n) espacio. La mejora real no es de complejidad sino de constantes: eliminaste el caso especial del for y la lógica de center, lo que simplifica el código pero no cambia la clase de complejidad. La versión canónica que existe para este problema también es O(n) / O(n). Este problema no tiene solución sub-lineal porque necesitas leer cada carácter al menos una vez:

def convert(self, s: str, numRows: int) -> str:
    if numRows == 1 or numRows >= len(s):
        return s

    rows = [''] * numRows
    row, direction = 0, -1

    for char in s:
        rows[row] += char
        if row == 0 or row == numRows - 1:
            direction *= -1
        row += direction

    return ''.join(rows)

Efectivamente: Claude me dio la mejor solución sin que se lo pidiera, y al analizarla pues tiene mucho sentido porque básicamente se trata de recorrer todas las letras mientras subes y bajas entre los renglones, cambiando la dirección a arriba cuando llegas al último renglón y cambiando la dirección a abajo cuando llegas al primero, y entonces en lugar de ver el zigzag así:

P A H N
APLSIIG
Y I R  

hay que verlo así:

P   A   H   N 
 A P L S I I G
  Y   I   R   

Esto me ayudó a entender que la solución es bastante obvia. ¿Por qué Claude me dio la solución así sin más? no sé, tal vez lo desesperé.

Problema anterior

Comentarios