cp_library/alg/iter/presum_fn.py

• View this file on GitHub
• Last update: 2025-03-30 20:17:47+09:00

Required by

• cp_library/alg/tree/fast/aux_tree_base_cls.py
• cp_library/alg/tree/fast/aux_tree_cls.py
• cp_library/alg/tree/fast/aux_tree_weighted_cls.py
• cp_library/alg/tree/lca_table_iterative_cls.py
• cp_library/alg/tree/lca_table_weighted_iterative_cls.py
• cp_library/alg/tree/tree_cls.py
• cp_library/alg/tree/tree_proto.py
• cp_library/alg/tree/tree_weighted_cls.py
• cp_library/alg/tree/tree_weighted_proto.py

Verified with

• test/aoj/grl/grl_5_a_diameter.test.py
• test/aoj/grl/grl_5_c_lca_table_iterative.test.py
• test/aoj/vol/0439_aux_dijkstra.test.py
• test/aoj/vol/0439_aux_rerooting_dp.test.py
• test/aoj/vol/0439_aux_weighted_rerooting_dp.test.py
• test/atcoder/abc/abc175_d_permutation.test.py
• test/atcoder/abc/abc202_e_dfs_enter_leave.test.py
• test/atcoder/abc/abc294_g_fast_tree_lca_table_weighted_bit.test.py
• test/atcoder/abc/abc337_g_tree_inversion_hld_bit.test.py
• test/atcoder/abc/abc337_g_tree_inversion_hld_fast.test.py
• test/atcoder/abc/abc361_e_tree_diameter.test.py
• test/atcoder/dp/dp_v_subtree_rerooting_iterative.test.py
• test/atcoder/dp/dp_v_subtree_rerooting_recursive.test.py
• test/library-checker/data-structure/static_range_sum.test.py
• test/library-checker/tree/jump_on_tree.test.py
• test/library-checker/tree/lca.test.py
• test/yukicoder/3407.test.py

Code

import cp_library.alg.iter.__header__
import operator
from itertools import accumulate
from typing import Callable, Iterable, TypeVar
from cp_library.misc.typing import _T

def presum(iter: Iterable[_T], func: Callable[[_T,_T],_T] = None, initial: _T = None, step = 1) -> list[_T]:
    if step == 1:
        return list(accumulate(iter, func, initial=initial))
    else:
        assert step >= 2
        if func is None:
            func = operator.add
        A = list(iter)
        if initial is not None:
            A = [initial] + A
        for i in range(step,len(A)):
            A[i] = func(A[i], A[i-step])
        return A

'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''
import operator
from itertools import accumulate
from typing import Callable, Iterable, TypeVar
_T = TypeVar('T')

def presum(iter: Iterable[_T], func: Callable[[_T,_T],_T] = None, initial: _T = None, step = 1) -> list[_T]:
    if step == 1:
        return list(accumulate(iter, func, initial=initial))
    else:
        assert step >= 2
        if func is None:
            func = operator.add
        A = list(iter)
        if initial is not None:
            A = [initial] + A
        for i in range(step,len(A)):
            A[i] = func(A[i], A[i-step])
        return A

Back to top page