ABC297 E — Kth Takoyaki Set

1# -*- coding: utf-8 -*-2 3 4import math5from bisect import bisect_left, bisect_right6from typing import Generic, Iterable, Iterator, TypeVar, Union, List7T = TypeVar('T')8 9 10class SortedSet(Generic[T]):11    """Sorted set (set) in C++.12    See:13    https://qiita.com/tatyam/items/492c70ac4c955c05560214    https://github.com/tatyam-prime/SortedSet/blob/main/SortedSet.py15    """16 17    BUCKET_RATIO = 5018    REBUILD_RATIO = 17019 20    def _build(self, a=None) -> None:21        "Evenly divide `a` into buckets."22        if a is None:23            a = list(self)24 25        size = self.size = len(a)26        bucket_size = int(math.ceil(math.sqrt(size / self.BUCKET_RATIO)))27        self.a = [a[size * i // bucket_size: size * (i + 1) // bucket_size] for i in range(bucket_size)]28 29    def __init__(self, a: Iterable[T] = []) -> None:30        """Make a new SortedSet from iterable.31            / O(N) if sorted and unique / O(N log N)32        """33        a = list(a)34 35        if not all(a[i] < a[i + 1] for i in range(len(a) - 1)):   # type: ignore36            a = sorted(set(a))  # type: ignore37 38        self._build(a)39 40    def __iter__(self) -> Iterator[T]:41        for i in self.a:42            for j in i:43                yield j  # type: ignore44 45    def __reversed__(self) -> Iterator[T]:46        for i in reversed(self.a):47            for j in reversed(i):48                yield j49 50    def __len__(self) -> int:51        return self.size52 53    def __repr__(self) -> str:54        return "SortedSet" + str(self.a)55 56    def __str__(self) -> str:57        s = str(list(self))58        return "{" + s[1: len(s) - 1] + "}"59 60    def _find_bucket(self, x: T) -> List[T]:61        "Find the bucket which should contain x. self must not be empty."62        for a in self.a:63            if x <= a[-1]:   # type: ignore64                return a65        return a66 67    def __contains__(self, x: T) -> bool:68        if self.size == 0:69            return False70 71        a = self._find_bucket(x)72        i = bisect_left(a, x)  # type: ignore73 74        return i != len(a) and a[i] == x75 76    def add(self, x: T) -> bool:77        "Add an element and return True if added. / O(√N)"78        if self.size == 0:79            self.a = [[x]]80            self.size = 181            return True82 83        a = self._find_bucket(x)84        i = bisect_left(a, x)  # type: ignore85 86        if i != len(a) and a[i] == x:87            return False88 89        a.insert(i, x)90        self.size += 191 92        if len(a) > len(self.a) * self.REBUILD_RATIO:93            self._build()94 95        return True96 97    def discard(self, x: T) -> bool:98        "Remove an element and return True if removed. / O(√N)"99        if self.size == 0:100            return False101 102        a = self._find_bucket(x)103        i = bisect_left(a, x)  # type: ignore104 105        if i == len(a) or a[i] != x:106            return False107 108        a.pop(i)109        self.size -= 1110 111        if len(a) == 0:112            self._build()113        return True114 115    def lt(self, x: T) -> Union[T, None]:116        "Find the largest element < x, or None if it doesn't exist."117        for a in reversed(self.a):118            if a[0] < x:  # type: ignore119                return a[bisect_left(a, x) - 1]  # type: ignore120        return None121 122    def le(self, x: T) -> Union[T, None]:123        "Find the largest element <= x, or None if it doesn't exist."124        for a in reversed(self.a):125            if a[0] <= x:  # type: ignore126                return a[bisect_right(a, x) - 1]  # type: ignore127        return None128 129    def gt(self, x: T) -> Union[T, None]:130        "Find the smallest element > x, or None if it doesn't exist."131        for a in self.a:132            if a[-1] > x:  # type: ignore133                return a[bisect_right(a, x)]  # type: ignore134        return None135 136    def ge(self, x: T) -> Union[T, None]:137        "Find the smallest element >= x, or None if it doesn't exist."138        for a in self.a:139            if a[-1] >= x:   # type: ignore140                return a[bisect_left(a, x)]  # type: ignore141        return None142 143    def __getitem__(self, x: int) -> T:144        "Return the x-th element, or IndexError if it doesn't exist."145        if x < 0:146            x += self.size147        if x < 0:148            raise IndexError149 150        for a in self.a:151            if x < len(a):152                return a[x]  # type: ignore153 154            x -= len(a)155        raise IndexError156 157    def index(self, x: T) -> int:158        "Count the number of elements < x."159        ans = 0160 161        for a in self.a:162            if a[-1] >= x:  # type: ignore163                return ans + bisect_left(a, x)  # type: ignore164            ans += len(a)165        return ans166 167    def index_right(self, x: T) -> int:168        "Count the number of elements <= x."169        ans = 0170 171        for a in self.a:172            if a[-1] > x:  # type: ignore173                return ans + bisect_right(a, x)  # type: ignore174            ans += len(a)175        return ans176 177 178def main():179    import sys180 181    input = sys.stdin.readline182 183    n, k = map(int, input().split())184    a = list(map(int, input().split()))185    s = SortedSet([0])186 187    for i in range(k):188        min_value = s[i]189 190        for ai in a:191            candidate = min_value + ai192 193            if candidate not in s:194                s.add(candidate)195    196    print(s[k])197 198 199if __name__ == "__main__":200    main()201