peterlu14
99 lines
2.0 KiB
Python
99 lines
2.0 KiB
Python
from collections import Counter as _Counter
|
|
from itertools import chain, islice
|
|
from typing import Tuple, Iterable, TypeVar, Callable
|
|
|
|
T = TypeVar('T')
|
|
|
|
try:
|
|
from typing import Counter
|
|
except ImportError:
|
|
# Counter is available in python 3.5.4, but we use 3.5.3 in Raspberry Pi.
|
|
from typing import Dict
|
|
|
|
Counter = Dict[T, int]
|
|
|
|
|
|
def count(it: Iterable[T]) -> int:
|
|
ret = 0
|
|
|
|
for _ in it:
|
|
ret += 1
|
|
|
|
return ret
|
|
|
|
|
|
def flat_iter(it: Iterable[Iterable[T]]) -> Iterable[T]:
|
|
return chain.from_iterable(it)
|
|
|
|
|
|
def flat_iter2(it: Iterable[Iterable[Iterable[T]]]) -> Iterable[T]:
|
|
return chain.from_iterable(chain.from_iterable(it))
|
|
|
|
|
|
def flat_iterN(it, level: int):
|
|
rt = it
|
|
|
|
while level > 0:
|
|
rt = chain.from_iterable(rt)
|
|
level -= 1
|
|
|
|
return rt
|
|
|
|
|
|
def zip_next2(it, level: int = 0):
|
|
it1 = flat_iterN(it, level)
|
|
it2 = flat_iterN(it, level)
|
|
|
|
return zip(it1, islice(it2, 1, None))
|
|
|
|
|
|
def zip_next3(it, level: int = 0):
|
|
it1 = flat_iterN(it, level)
|
|
it2 = flat_iterN(it, level)
|
|
it3 = flat_iterN(it, level)
|
|
|
|
return zip(it1, islice(it2, 1, None), islice(it3, 2, None))
|
|
|
|
|
|
def counter_map(s: Iterable[Tuple[T, int]]) -> Counter[T]:
|
|
c = _Counter()
|
|
|
|
for k, v in s:
|
|
c[k] += v
|
|
|
|
return c
|
|
|
|
|
|
def counter_total(c: Counter) -> int:
|
|
return sum(c.values())
|
|
|
|
|
|
def counter_product_sum(c: Counter[T],
|
|
key_mapper: Callable[[T], int] = None,
|
|
value_mapper: Callable[[int], int] = None) -> int:
|
|
s = 0
|
|
for k, v in c.items():
|
|
_k = k if key_mapper is None else key_mapper(k)
|
|
_v = v if value_mapper is None else value_mapper(v)
|
|
s += _k * _v
|
|
return s
|
|
|
|
|
|
def counter_diff(v: Tuple[int, int]) -> int:
|
|
return v[1] - v[0]
|
|
|
|
|
|
def counter_miss(v1: int, v2: int, v3: int) -> int:
|
|
if (v2 - v1) == (v2 - v2):
|
|
return 0
|
|
else:
|
|
return v3 - v1
|
|
|
|
|
|
def counter_non_zero(c: Counter[T]) -> Counter[T]:
|
|
for k in list(c.keys()):
|
|
if c[k] == 0:
|
|
del c[k]
|
|
|
|
return c
|