improved implementation of rolling_partition.

This is about 20% faster than the old version by the benchmark in
bench/test_rolling_partition.py

Also added imports

Implementation by @eriknw