Coverage for src\aoc2022\days\day9.py: 100%

1import math 

2from enum import Enum 

3from typing import List, Set, Tuple 

4 

5import numpy as np 

6 

7from src.aoc2022 import utils 

8 

9 

10class Move(Enum): 

11 """ 

12 Enum for valid rope moves 

13 """ 

14 

15 LEFT = "L" 

16 """ 

17 Value for 'LEFT' move. 

18 """ 

19 RIGHT = "R" 

20 """ 

21 Value for 'RIGHT' move. 

22 """ 

23 UP = "U" 

24 """ 

25 Value for 'UP' move. 

26 """ 

27 DOWN = "D" 

28 """ 

29 Value for 'DOWN' move. 

30 """ 

31 

32 

33def parse_data_to_array(raw_data: str) -> List[Tuple[Move, int]]: 

34 """ 

35 Parse the raw data into a 1D list of move Tuples 

36 

37 :param raw_data: Raw data containing numerical data 

38 :type raw_data: str 

39 

40 :return: parsed matrix of values 

41 :rtype: List[Tuple[Move, int]] 

42 """ 

43 parsed_list = [] 

44 for line in raw_data.split("\n"): 

45 move, count = line.split(" ") 

46 parsed_list.append((Move(move), int(count))) 

47 

48 return parsed_list 

49 

50 

51def get_tail_update( 

52 head_pos: Tuple[int, int], tail_pos: Tuple[int, int] 

53) -> Tuple[int, int]: 

54 

55 x1, y1 = head_pos 

56 x2, y2 = tail_pos 

57 

58 L2_distance = math.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2) 

59 

60 if L2_distance >= 2: 

61 x_move = np.sign(x1 - x2) 

62 y_move = np.sign(y1 - y2) 

63 

64 return (x2 + x_move, y2 + y_move) 

65 else: 

66 return (x2, y2) 

67 

68 

69def get_head_update(pos: Tuple[int, int], move: Move) -> Tuple[int, int]: 

70 x, y = pos 

71 

72 if move == Move.UP: 

73 return (x, y + 1) 

74 elif move == Move.DOWN: 

75 return (x, y - 1) 

76 elif move == Move.RIGHT: 

77 return (x + 1, y) 

78 else: 

79 return (x - 1, y) 

80 

81 

82def get_knot_positions(moves: List[Tuple[Move, int]]) -> Set[Tuple[int, int]]: 

83 head_pos = (0, 0) 

84 tail_pos = (0, 0) 

85 unique_positions = set() 

86 

87 for move, count in moves: 

88 for i in range(count): 

89 head_pos = get_head_update(head_pos, move) 

90 tail_pos = get_tail_update(head_pos, tail_pos) 

91 unique_positions.add(tail_pos) 

92 

93 return unique_positions 

94 

95 

96def get_multi_knot_positions( 

97 moves: List[Tuple[Move, int]], num_knots: int 

98) -> Set[Tuple[int, int]]: 

99 unique_positions = set() 

100 knot_positions = [(0, 0) for _ in range(num_knots)] 

101 for move, count in moves: 

102 for i in range(count): 

103 

104 head_pos = knot_positions[0] 

105 knot_positions[0] = get_head_update(head_pos, move) 

106 

107 for j in range(num_knots - 1): 

108 head_pos = knot_positions[j] 

109 tail_pos = knot_positions[j + 1] 

110 knot_positions[j + 1] = get_tail_update(head_pos, tail_pos) 

111 

112 unique_positions.add(knot_positions[-1]) 

113 

114 return unique_positions 

115 

116 

117if __name__ == "__main__": # pragma: no cover 

118 

119 data = utils.get_raw_data("./src/aoc2022/data/day9.txt") 

120 

121 parsed_array = parse_data_to_array(data) 

122 unique_positions = get_knot_positions(parsed_array) 

123 

124 print(f"{len(unique_positions)} unique positions were visited") 

125 

126 unique_positions = get_multi_knot_positions(parsed_array, 10) 

127 

128 print(f"{len(unique_positions)} unique positions were visited")