Coverage for blog/dsa/leetcode/zigzag/__init__.py: 53%

1import pytest 

2 

3 

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 

12 

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) 

16 

17 for kk in range(1, n - 1): 

18 

19 for jj in range(jj_max): 

20 

21 aa = (incr * jj) + kk 

22 bb = (incr * (jj + 1)) - kk 

23 

24 if aa < m: 

25 yield aa 

26 if bb < m: 

27 yield bb 

28 

29 yield from (num for jj in range(jj_max) if (num := (incr * jj) + N) < m) 

30 

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)) 

35 

36 

37# end snippet solution 

38 

39 

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)] 

46 

47 current_row = 0 

48 incr = -1 

49 

50 for char in s: 

51 if current_row == N or current_row == 0: 

52 incr = -1 * incr 

53 

54 rows[current_row] += char 

55 current_row = current_row + incr 

56 

57 return "".join(rows) 

58 

59 

60# end snippet solution_better 

61 

62 

63@pytest.fixture 

64def solution(): 

65 return SolutionBetter() 

66 

67 

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): 

86 

87 assert solution.convert(question, numRows) == answer