Leetcode 6. ZigZag Conversion with Python - 리뷰 O

문제

The string "PAYPALISHIRING" is written in a zigzag pattern on a given number of rows like this: 

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

And then read line by line: "PAHNAPLSIIGYIR"

주어진 문자열에 대해 위와 같이 지그재그로 나열해 새로 읽은 문자열로 변환해 리턴하시오.

예시

Input: s = "PAYPALISHIRING", numRows = 3
Output: "PAHNAPLSIIGYIR"

Input: s = "PAYPALISHIRING", numRows = 4
Output: "PINALSIGYAHRPI"
Explanation:
P     I    N
A   L S  I G
Y A   H R
P     I

제한

• 1 <= s.length <= 1000
• s 는 영문자와 ',' 및 '.' 으로만 이루어져 있다.
• 1 <= numRows <= 1000

풀이

각각의 row에 대해 포함 여부를 index로 계산해 채워 나간다.

class Solution:
    def convert(self, s: str, numRows: int) -> str:
        if numRows == 1: return s
        group = (numRows-1)*2
        ans = ''
        for j in range(0, numRows):
            for i, v in enumerate(s):
                if i % group == j or i % group == group - j:
                    ans += v
        return ans
        

리뷰

• 시간복잡도의 큰 개선

class Solution:
    def convert(self, s: str, numRows: int) -> str:

        if numRows == 1:
            return s

        group = (numRows - 1) * 2
        ans = [[] for _ in range(numRows)]

        for i, ch in enumerate(s):
            place = i % group

            if place < numRows:
                ans[place].append(ch)
            else:
                ans[group - place].append(ch)

        return "".join("".join(row) for row in ans)