Coverage for blog/dsa/leetcode/zigzag/__init__.py: 53%
43 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-20 16:23 +0000
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-20 16:23 +0000
1import pytest
4# start snippet solution
5class Solution:
6 def coords(self, m: int, n_rows: int):
7 n = n_rows
8 N = n - 1
9 if N == 0:
10 yield from range(m)
11 return
13 incr = 2 * N
14 jj_max = (m // incr) + 1
15 yield from (num for jj in range(jj_max) if (num := incr * jj) < m)
17 for kk in range(1, n - 1):
19 for jj in range(jj_max):
21 aa = (incr * jj) + kk
22 bb = (incr * (jj + 1)) - kk
24 if aa < m:
25 yield aa
26 if bb < m:
27 yield bb
29 yield from (num for jj in range(jj_max) if (num := (incr * jj) + N) < m)
31 def convert(self, s: str, numRows: int) -> str:
32 m = len(s)
33 nums = self.coords(m, numRows)
34 return "".join((s[num] for num in nums))
37# end snippet solution
40# start snippet solution_better
41class SolutionBetter:
42 def convert(self, s: str, numRows: int) -> str:
43 n = len(s)
44 N = numRows - 1
45 rows = ["" for _ in range(n)]
47 current_row = 0
48 incr = -1
50 for char in s:
51 if current_row == N or current_row == 0:
52 incr = -1 * incr
54 rows[current_row] += char
55 current_row = current_row + incr
57 return "".join(rows)
60# end snippet solution_better
63@pytest.fixture
64def solution():
65 return SolutionBetter()
68@pytest.mark.parametrize(
69 "question, numRows, answer",
70 (
71 # From post
72 ("AAAABBCCCCDD", 4, "ACABCDABCDAC"),
73 ("PAYPALISHIRING", 3, "PAHNAPLSIIGYIR"),
74 ("PAYPALISHIRING", 4, "PINALSIGYAHRPI"),
75 # Trivial
76 ("A", 5, "A"),
77 ("A", 1, "A"),
78 ("ABCDE", 4, "ABCED"),
79 # a e i
80 # b d f h j l
81 # c g k
82 ("abcdefghijkl", 3, "aeibdfhjlcgk"),
83 ),
84)
85def test_solution(solution: Solution, question: str, numRows: int, answer: str):
87 assert solution.convert(question, numRows) == answer