1
- from typing import Tuple
1
+ from typing import (
2
+ Optional ,
3
+ Sequence ,
4
+ Tuple ,
5
+ Union
6
+ )
2
7
from datetime import date , timedelta
3
8
from epiweeks import Week , Year
4
9
@@ -23,6 +28,10 @@ def guess_time_value_is_day(value: int) -> bool:
23
28
# YYYYMMDD type and not YYYYMM
24
29
return len (str (value )) > 6
25
30
31
+ def guess_time_value_is_week (value : int ) -> bool :
32
+ # YYYYWW type and not YYYYMMDD
33
+ return len (str (value )) == 6
34
+
26
35
def date_to_time_value (d : date ) -> int :
27
36
return int (d .strftime ("%Y%m%d" ))
28
37
@@ -67,3 +76,95 @@ def weeks_in_range(week_range: Tuple[int, int]) -> int:
67
76
year = Year (y )
68
77
acc += year .totalweeks ()
69
78
return acc + 1 # same week should lead to 1 week that will be queried
79
+
80
+ def dates_to_ranges (values : Optional [Sequence [Union [Tuple [int , int ], int ]]]) -> Optional [Sequence [Union [Tuple [int , int ], int ]]]:
81
+ """
82
+ Converts a mixed list of dates and date ranges to an optimized list where dates are merged into ranges where possible.
83
+ e.g. [20200101, 20200102, (20200101, 20200104), 20200106] -> [(20200101, 20200104), 20200106]
84
+ (the first two values of the original list are merged into a single range)
85
+ """
86
+ if not values or len (values ) <= 1 :
87
+ return values
88
+
89
+ # determine whether the list is of days (YYYYMMDD) or weeks (YYYYWW) based on first element
90
+ try :
91
+ if (isinstance (values [0 ], tuple ) and guess_time_value_is_day (values [0 ][0 ]))\
92
+ or (isinstance (values [0 ], int ) and guess_time_value_is_day (values [0 ])):
93
+ return days_to_ranges (values )
94
+ elif (isinstance (values [0 ], tuple ) and guess_time_value_is_week (values [0 ][0 ]))\
95
+ or (isinstance (values [0 ], int ) and guess_time_value_is_week (values [0 ])):
96
+ return weeks_to_ranges (values )
97
+ else :
98
+ return values
99
+ except :
100
+ return values
101
+
102
+ def days_to_ranges (values : Sequence [Union [Tuple [int , int ], int ]]) -> Sequence [Union [Tuple [int , int ], int ]]:
103
+ intervals = []
104
+
105
+ # populate list of intervals based on original values
106
+ for v in values :
107
+ if isinstance (v , int ):
108
+ # 20200101 -> [20200101, 20200101]
109
+ intervals .append ([time_value_to_date (v ), time_value_to_date (v )])
110
+ else : # tuple
111
+ # (20200101, 20200102) -> [20200101, 20200102]
112
+ intervals .append ([time_value_to_date (v [0 ]), time_value_to_date (v [1 ])])
113
+
114
+ intervals .sort (key = lambda x : x [0 ])
115
+
116
+ # merge overlapping intervals https://leetcode.com/problems/merge-intervals/
117
+ merged = []
118
+ for interval in intervals :
119
+ # no overlap, append the interval
120
+ # caveat: we subtract 1 from interval[0] so that contiguous intervals are considered overlapping. i.e. [1, 1], [2, 2] -> [1, 2]
121
+ if not merged or merged [- 1 ][1 ] < interval [0 ] - timedelta (days = 1 ):
122
+ merged .append (interval )
123
+ # overlap, merge the current and previous intervals
124
+ else :
125
+ merged [- 1 ][1 ] = max (merged [- 1 ][1 ], interval [1 ])
126
+
127
+ # convert intervals from dates back to integers
128
+ ranges = []
129
+ for m in merged :
130
+ if m [0 ] == m [1 ]:
131
+ ranges .append (date_to_time_value (m [0 ]))
132
+ else :
133
+ ranges .append ((date_to_time_value (m [0 ]), date_to_time_value (m [1 ])))
134
+
135
+ return ranges
136
+
137
+ def weeks_to_ranges (values : Sequence [Union [Tuple [int , int ], int ]]) -> Sequence [Union [Tuple [int , int ], int ]]:
138
+ intervals = []
139
+
140
+ # populate list of intervals based on original values
141
+ for v in values :
142
+ if isinstance (v , int ):
143
+ # 202001 -> [202001, 202001]
144
+ intervals .append ([week_value_to_week (v ), week_value_to_week (v )])
145
+ else : # tuple
146
+ # (202001, 202002) -> [202001, 202002]
147
+ intervals .append ([week_value_to_week (v [0 ]), week_value_to_week (v [1 ])])
148
+
149
+ intervals .sort (key = lambda x : x [0 ])
150
+
151
+ # merge overlapping intervals https://leetcode.com/problems/merge-intervals/
152
+ merged = []
153
+ for interval in intervals :
154
+ # no overlap, append the interval
155
+ # caveat: we subtract 1 from interval[0] so that contiguous intervals are considered overlapping. i.e. [1, 1], [2, 2] -> [1, 2]
156
+ if not merged or merged [- 1 ][1 ] < interval [0 ] - 1 :
157
+ merged .append (interval )
158
+ # overlap, merge the current and previous intervals
159
+ else :
160
+ merged [- 1 ][1 ] = max (merged [- 1 ][1 ], interval [1 ])
161
+
162
+ # convert intervals from weeks back to integers
163
+ ranges = []
164
+ for m in merged :
165
+ if m [0 ] == m [1 ]:
166
+ ranges .append (week_to_time_value (m [0 ]))
167
+ else :
168
+ ranges .append ((week_to_time_value (m [0 ]), week_to_time_value (m [1 ])))
169
+
170
+ return ranges
0 commit comments