Problem A2: Perfectly Balanced - Chapter 2

meta_hackercup_a2_perfectly_balanced.rs

/*
https://www.facebook.com/codingcompetitions/hacker-cup/2022/round-2/problems/A2
https://youtu.be/CsSGu51V7YQ
*/

use std::collections::HashSet;
use rand::{thread_rng, Rng};
use std::io;
use std::io::Read;
use crate::fenwick_tree::FenwickTree;

fn get_nums() -> Vec<usize> {
    let mut buf = vec![];
    let mut nums = vec![];
    io::stdin().read_to_end(&mut buf).unwrap();
    let mut cur = 0;
    for b in buf {
        if b >= b'0' && b <= b'9' {
            cur = cur * 10 + (b - b'0') as usize;
        } else {
            if cur > 0 {
                nums.push(cur);
            }
            cur = 0;
        }
    }
    nums
}

pub mod fenwick_tree {
    /// `FenwickTree` is a data structure that can efficiently update elements
    /// and calculate prefix sums in a table of numbers.
    /// [https://en.wikipedia.org/wiki/Fenwick_tree](https://en.wikipedia.org/wiki/Fenwick_tree)
    pub struct FenwickTree<T, F> {
        n: usize,
        data: Vec<T>,
        initialize: F,
    }

    impl<T, F> FenwickTree<T, F>
        where
            T: Copy + std::ops::AddAssign + std::ops::Sub<Output = T>,
            F: Fn() -> T,
    {
        /// Constructs a new `FenwickTree`. The size of `FenwickTree` should be specified by `size`.
        pub fn new(size: usize, initialize: F) -> FenwickTree<T, F> {
            FenwickTree {
                n: size + 1,
                data: vec![initialize(); size + 1],
                initialize,
            }
        }

        pub fn add(&mut self, k: usize, value: T) {
            let mut x = k;
            while x < self.n {
                self.data[x] += value;
                x |= x + 1;
            }
        }

        /// Returns a sum of range `[l, r)`
        pub fn sum(&self, l: usize, r: usize) -> T {
            self.sum_one(r) - self.sum_one(l)
        }

        /// Returns a sum of range `[0, k)`
        pub fn sum_one(&self, k: usize) -> T {
            assert!(k < self.n, "Cannot calculate for range [{}, {})", k, self.n);
            let mut result = (self.initialize)();
            let mut x = k as i32 - 1;
            while x >= 0 {
                result += self.data[x as usize];
                x = (x & (x + 1)) - 1;
            }

            result
        }
    }
}


fn main() {
    let mut it = get_nums().into_iter();
    let mut g = || it.next().unwrap();
    let mut rand = vec![0; 10usize.pow(6) + 5];
    let mut rng = thread_rng();
    for r in &mut rand {
        *r = rng.gen_range(0..10i64.pow(13));
    }
    let hash_set: HashSet<i64> = rand.iter().cloned().collect();
    let t = g();
    for test_case in 1..=t {
        let n = g();
        let mut a = vec![0; n];

        let mut t = FenwickTree::new(n, || 0);
        for (i, v) in a.iter_mut().enumerate() {
            *v = g();
            t.add(i, rand[*v]);
        }
        let q = g();
        let mut ans = 0;
        for _ in 0..q {
            let c = g();
            if c == 1 {
                let (mut x, y) = (g(), g());
                x -= 1;
                t.add(x, rand[y]-rand[a[x]]);
                a[x] = y;
            } else {
                let (mut l, r) = (g(), g());
                l -= 1;
                let sz = r - l;
                if sz % 2 == 0 {
                    continue;
                }
                for mid in l+sz/2..=l+sz/2+1 {
                    let a = t.sum(l, mid);
                    let b = t.sum(mid, r);
                    let diff = a.max(b) - a.min(b);
                    if hash_set.contains(&diff) {
                        ans += 1;
                        break;
                    }
                }
            }
        }
        println!("Case #{}: {}", test_case, ans);
    }
}


/*

     author:  tourist
     created: 24.09.2022 20:00:50       
#undef _GLIBCXX_DEBUG

#include <bits/stdc++.h>

using namespace std;

#ifdef LOCAL
#include "algo/debug.h"
#else
#define debug(...) 42
#endif

template <typename T>
class fenwick {
 public:
  vector<T> fenw;
  int n;

  fenwick(int _n) : n(_n) {
    fenw.resize(n);
  }

  void modify(int x, T v) {
    while (x < n) {
      fenw[x] += v;
      x |= (x + 1);
    }
  }

  T get(int x) {
    T v{};
    while (x >= 0) {
      v += fenw[x];
      x = (x & (x + 1)) - 1;
    }
    return v;
  }
};

typedef unsigned long long ull;

mt19937_64 rng((unsigned int) chrono::steady_clock::now().time_since_epoch().count());

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  const int M = (int) 1e6 + 10;
  vector<ull> coeff(M);
  for (int i = 0; i < M; i++) {
    coeff[i] = rng();
  }
  auto all = coeff;
  sort(all.begin(), all.end());
  int tt;
  cin >> tt;
  for (int qq = 1; qq <= tt; qq++) {
    cout << "Case #" << qq << ": ";
    int n;
    cin >> n;
    vector<int> s(n);
    for (int i = 0; i < n; i++) {
      cin >> s[i];
    }
    fenwick<ull> fenw(n);
    for (int i = 0; i < n; i++) {
      fenw.modify(i, coeff[s[i]]);
    }
    int ans = 0;
    int q;
    cin >> q;
    while (q--) {
      int op;
      cin >> op;
      if (op == 2) {
        int l, r;
        cin >> l >> r;
        --l;
        int len = r - l;
        if (len % 2 == 1) {
          bool ok = false;
          for (int x = l + len / 2; x <= l + len / 2 + 1; x++) {
            auto a = fenw.get(x - 1) - fenw.get(l - 1);
            auto b = fenw.get(r - 1) - fenw.get(x - 1);
            auto diff = (x == l + len / 2 ? b - a : a - b);
            auto it = lower_bound(all.begin(), all.end(), diff);
            if (it != all.end() && (*it) == diff) {
              ok = true;
              break;
            }
          }
          if (ok) {
            ans += 1;
          }
        }
      } else {
        int x, y;
        cin >> x >> y;
        --x;
        fenw.modify(x, coeff[y] - coeff[s[x]]);
        s[x] = y;
      }
    }
    cout << ans << '\n';
  }
  return 0;
}

*/