You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

61 lines
1.6 KiB

4 years ago
  1. from cpython.sequence cimport PySequence_Tuple
  2. from cytoolz.itertoolz cimport frequencies, pluck
  3. from itertools import groupby
  4. from cytoolz.compatibility import map
  5. __all__ = ['countby', 'partitionby']
  6. cpdef object countby(object key, object seq):
  7. """
  8. Count elements of a collection by a key function
  9. >>> countby(len, ['cat', 'mouse', 'dog'])
  10. {3: 2, 5: 1}
  11. >>> def iseven(x): return x % 2 == 0
  12. >>> countby(iseven, [1, 2, 3]) # doctest:+SKIP
  13. {True: 1, False: 2}
  14. See Also:
  15. groupby
  16. """
  17. if not callable(key):
  18. return frequencies(pluck(key, seq))
  19. return frequencies(map(key, seq))
  20. cdef class partitionby:
  21. """ partitionby(func, seq)
  22. Partition a sequence according to a function
  23. Partition `s` into a sequence of lists such that, when traversing
  24. `s`, every time the output of `func` changes a new list is started
  25. and that and subsequent items are collected into that list.
  26. >>> is_space = lambda c: c == " "
  27. >>> list(partitionby(is_space, "I have space"))
  28. [('I',), (' ',), ('h', 'a', 'v', 'e'), (' ',), ('s', 'p', 'a', 'c', 'e')]
  29. >>> is_large = lambda x: x > 10
  30. >>> list(partitionby(is_large, [1, 2, 1, 99, 88, 33, 99, -1, 5]))
  31. [(1, 2, 1), (99, 88, 33, 99), (-1, 5)]
  32. See also:
  33. partition
  34. groupby
  35. itertools.groupby
  36. """
  37. def __cinit__(self, object func, object seq):
  38. self.iter_groupby = groupby(seq, key=func)
  39. def __iter__(self):
  40. return self
  41. def __next__(self):
  42. cdef object key, val
  43. key, val = next(self.iter_groupby)
  44. return PySequence_Tuple(val)