NEPLg2 Tutorial - 25 - prefix sum と two pointers
Web Playground
Web Playground

prefix sum と two pointers

O(N^2) を O(N) へ落とす定番として、累積和と尺取り法を扱います。
この章は Vec と std/streamio / kp/* の補助 API を優先して、手書きメモリ操作を減らします。

prefix sum で区間和を O(1) にする

sum[l..r) = pref[r] - pref[l] を使います。

TESTstdionormalize_newlines
#entry main
#indent 4
#target std

#import "core/math" as *
#import "core/result" as *
#import "std/streamio" as *
#import "std/iotarget" as *
#import "kp/kpprefix" as *
#import "alloc/collections/vec" as *

fn main <()*> ()> ():
    let sc <StreamScanner> unwrap_ok open ReadStream::Stdio;
    let n <i32> read sc;
    let q <i32> read sc;
    let mut a <Vec<i32>> new<i32>;
    let mut i <i32> 0;
    while lt i n:
        do:
            set a push a read sc;
            set i add i 1;
    let pref <PrefixI32> prefix_build_vec_i32 a;
    let mut w <StreamWriter> unwrap_ok open WriteStream::Stdio;
    let mut k <i32> 0;
    while lt k q:
        do:
            let l1 <i32> read sc;
            let r1 <i32> read sc;
            let ans <i32> prefix_sum_i32 pref sub l1 1 r1;
            set w writeln w ans;
            set k add k 1;
    close sc;
    set w flush w;
    close w;
    prefix_free_i32 pref
stdin
5 3
1 2 3 4 5
1 3
2 5
1 5
stdout
6
14
15

two pointers で sum <= S の部分配列数を数える

正の配列なら、右端を戻さない 2 ポインタで O(N) にできます。

TESTstdionormalize_newlines
#entry main
#indent 4
#target std

#import "core/math" as *
#import "core/option" as *
#import "alloc/collections/vec" as *
#import "std/stdio" as *

fn count_subarrays_leq_s <(Vec<i32>,i32)->i32> (a, s):
    let n <i32> vec_len<i32> a;
    let mut l <i32> 0;
    let mut r <i32> 0;
    let mut sum <i32> 0;
    let mut ans <i32> 0;
    while lt l n:
        do:
            let mut can_extend <i32> 1;
            while eq can_extend 1:
                do:
                    if lt r n:
                        then:
                            let rv <i32> unwrap<i32> vec_get<i32> a r;
                            if le add sum rv s:
                                then:
                                    set sum add sum rv;
                                    set r add r 1;
                                else set can_extend 0
                        else set can_extend 0
            set ans add ans sub r l;
            if lt l r:
                then set sum sub sum unwrap<i32> vec_get<i32> a l
                else set r add l 1
            set l add l 1;
    ans

fn main <()*> ()> ():
    let a <Vec<i32>>:
        new<i32>
        |> push 1
        |> push 2
        |> push 3
        |> push 4
    println_i32 count_subarrays_leq_s a 5
stdout
6