tipstar0125 / atcoder

atcoder
6 stars 0 forks source link

セグ木 #42

Open tipstar0125 opened 8 months ago

tipstar0125 commented 8 months ago
#[derive(Debug, Clone)]
struct SegmentTree<M>
where
    M: Monoid,
{
    n: usize,
    offset: usize,
    data: Vec<M::S>,
}

impl<M> SegmentTree<M>
where
    M: Monoid,
    M::S: Clone,
{
    fn new(n: usize) -> Self {
        let offset = n.next_power_of_two();
        let data = vec![M::id(); offset * 2];
        SegmentTree { n, offset, data }
    }
    // pos: 0-indexed!!
    fn set(&mut self, pos: usize, x: M::S) {
        let mut pos = pos + self.offset;
        self.data[pos] = x;
        while pos > 0 {
            pos /= 2;
            self.data[pos] = M::op(&self.data[pos * 2], &self.data[pos * 2 + 1]);
        }
    }
    // pos: 0-indexed!!
    fn get(&self, pos: usize) -> M::S {
        self.data[self.offset + pos].clone()
    }
    // 0-indexed!!, [l, r)
    fn prod(&self, mut l: usize, mut r: usize) -> M::S {
        l += self.offset;
        r += self.offset;
        let mut xl = M::id();
        let mut xr = M::id();

        while l < r {
            if l % 2 == 1 {
                xl = M::op(&xl, &self.data[l]);
                l += 1;
            }
            if r % 2 == 1 {
                r -= 1;
                xr = M::op(&xr, &self.data[r]);
            }
            l /= 2;
            r /= 2;
        }
        M::op(&xl, &xr)
    }
    fn all_prod(&self) -> M::S {
        self.data[1].clone()
    }
    fn max_right<F>(&self, mut l: usize, f: F) -> usize
    where
        F: Fn(&M::S) -> bool,
    {
        assert!(l <= self.n);
        assert!(f(&M::id()));
        if l == self.n {
            return self.n;
        }
        l += self.offset;
        let mut sm = M::id();
        loop {
            while l % 2 == 0 {
                l /= 2;
            }
            if !f(&M::op(&sm, &self.data[l])) {
                while l < self.offset {
                    l *= 2;
                    let res = M::op(&sm, &self.data[l]);
                    if f(&res) {
                        sm = res;
                        l += 1;
                    }
                }
                return l - self.offset;
            }
            sm = M::op(&sm, &self.data[l]);
            l += 1;
            if (l & l.wrapping_neg()) == l {
                break;
            }
        }
        self.n
    }
    fn min_left<F>(&self, mut r: usize, f: F) -> usize
    where
        F: Fn(&M::S) -> bool,
    {
        assert!(r <= self.n);
        assert!(f(&M::id()));
        if r == 0 {
            return 0;
        }
        r += self.offset;
        let mut sm = M::id();
        loop {
            r -= 1;
            while r > 1 && r % 2 == 1 {
                r >>= 1;
            }
            if !f(&M::op(&self.data[r], &sm)) {
                while r < self.offset {
                    r = 2 * r + 1;
                    let res = M::op(&self.data[r], &sm);
                    if f(&res) {
                        sm = res;
                        r -= 1;
                    }
                }
                return r + 1 - self.offset;
            }
            sm = M::op(&self.data[r], &sm);
            if (r & r.wrapping_neg()) == r {
                break;
            }
        }
        0
    }
}

trait Monoid {
    type S;
    fn op(a: &Self::S, b: &Self::S) -> Self::S;
    fn id() -> Self::S;
}

struct MinMonoid;

impl Monoid for MinMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        std::cmp::min(*a, *b)
    }
    fn id() -> Self::S {
        std::usize::MAX
    }
}

struct MaxMonoid;

impl Monoid for MaxMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        std::cmp::max(*a, *b)
    }
    fn id() -> Self::S {
        0
    }
}

struct AddMonoid;

impl Monoid for AddMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        *a + *b
    }
    fn id() -> Self::S {
        0
    }
}

struct MulMonoid;

impl Monoid for MulMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        *a * *b
    }
    fn id() -> Self::S {
        1
    }
}

struct GcdMonoid;

impl Monoid for GcdMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        let mut a = *a;
        let mut b = *b;
        while b != 0 {
            let q = a / b;
            let r = a - b * q;
            a = b;
            b = r;
        }
        a
    }
    fn id() -> Self::S {
        0
    }
}
tipstar0125 commented 8 months ago
tipstar0125 commented 8 months ago

0-indexed計算イメージ

image

image

image

image

tipstar0125 commented 8 months ago

不採用!! データ構造を0-indexedで構築。 データ構造が1-indexedのときの、メソッド(set, prod)のindex指定が1-indexedになると思ったが、そうではないので、0-indexedのデータ構造よりも、1-indexedのデータ構造の方が良さそう。

#[derive(Debug, Clone)]
struct SegmentTree<M>
where
    M: Monoid,
{
    offset: usize,
    data: Vec<M::S>,
}

impl<M> SegmentTree<M>
where
    M: Monoid,
    M::S: Clone,
{
    fn new(n: usize) -> Self {
        let offset = n.next_power_of_two();
        let data = vec![M::id(); offset * 2 - 1];
        SegmentTree { offset, data }
    }
    // pos: 0-indexed!!
    fn update(&mut self, pos: usize, x: M::S) {
        let mut pos = pos + self.offset - 1;
        self.data[pos] = x;
        while pos > 0 {
            pos = (pos - 1) / 2;
            self.data[pos] = M::op(&self.data[pos * 2 + 1], &self.data[pos * 2 + 2]);
        }
    }
    // pos: 0-indexed!!
    fn get(&self, pos: usize) -> M::S {
        self.data[self.offset + pos - 1].clone()
    }
    // 0-indexed!!, [l, r)
    fn query(&self, mut l: usize, mut r: usize) -> M::S {
        l += self.offset - 1;
        r += self.offset - 1;
        let mut xl = M::id();
        let mut xr = M::id();

        while l < r {
            if l % 2 == 0 {
                xl = M::op(&xl, &self.data[l]);
            }
            if r % 2 == 0 {
                xr = M::op(&xr, &self.data[r - 1]);
            }
            l /= 2;
            r = (r - 1) / 2;
        }
        M::op(&xl, &xr)
    }
}

trait Monoid {
    type S;
    fn op(a: &Self::S, b: &Self::S) -> Self::S;
    fn id() -> Self::S;
}

struct MinMonoid;

impl Monoid for MinMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        std::cmp::min(*a, *b)
    }
    fn id() -> Self::S {
        std::usize::MAX
    }
}

struct MaxMonoid;

impl Monoid for MaxMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        std::cmp::max(*a, *b)
    }
    fn id() -> Self::S {
        0
    }
}

struct AddMonoid;

impl Monoid for AddMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        *a + *b
    }
    fn id() -> Self::S {
        0
    }
}

struct MulMonoid;

impl Monoid for MulMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        *a * *b
    }
    fn id() -> Self::S {
        1
    }
}

struct GcdMonoid;

impl Monoid for GcdMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        let mut a = *a;
        let mut b = *b;
        while b != 0 {
            let q = a / b;
            let r = a - b * q;
            a = b;
            b = r;
        }
        a
    }
    fn id() -> Self::S {
        0
    }
}