Coverage for src / competitive_verifier / verify / split_state.py: 100%

1from typing import TypeVar 

2 

3from pydantic import BaseModel 

4 

5T = TypeVar("T") 

6 

7 

8class SplitState(BaseModel): 

9 size: int 

10 index: int 

11 

12 def __str__(self) -> str: 

13 return f"{self.index}/{self.size}" 

14 

15 def split(self, lst: list[T]) -> list[T]: 

16 """Split list. 

17 

18 Args: 

19 lst (list[T]): Target list 

20 

21 Returns: 

22 list[T]: Splited list 

23 

24 Example: 

25 state = SplitState(size=3, index=0) 

26 assert state.split([0, 1, 2, 3, 4]) == [0] 

27 state = SplitState(size=3, index=1) 

28 assert state.split([0, 1, 2, 3, 4]) == [1, 2] 

29 state = SplitState(size=3, index=2) 

30 assert state.split([0, 1, 2, 3, 4]) == [3, 4] 

31 state = SplitState(size=6, index=4) 

32 assert state.split([0, 1, 2, 3, 4]) == [4] 

33 state = SplitState(size=6, index=5) 

34 assert state.split([0, 1, 2, 3, 4]) == [] 

35 """ 

36 if len(lst) <= self.size: 

37 if len(lst) <= self.index: 

38 return [] 

39 return [lst[self.index]] 

40 

41 from_index = len(lst) * self.index // self.size 

42 to_index = len(lst) * (self.index + 1) // self.size 

43 return lst[from_index:to_index]