HashMapでdpすることができなくもない

概要

  • RustのHashMapのラッパーを作ってIndex, IndexMutトレイトを実装すると配列のように使えて便利
  • でも流石に大きな定義域でループを回すとTLEする

HashMapをindexで参照と可変参照できるようにする

github.com

昔はあったらしいHashMapのIndexMut実装は↑によって消されてしまったので独自実装する。

#[allow(unused_macros)]
macro_rules! map {
    () => {
        IndexMap(std::collections::HashMap::new())
    };
    ($value:expr; $key_iter:expr) => {
            IndexMap(($key_iter).map(|k| (k, $value)).collect::<std::collections::HashMap<_,_>>())
    };
    ($value:expr; $($key:expr);*) => {
        map!($($key);* => $value)
    };
    ($head:expr; $($tail:expr);* => $init:expr) => {
        map![map!($($tail);* => $init); $head]
    };
    ($head:expr => $init:expr) => {
        map![$init; $head]
    };
}

// HashMapのラッパー
#[derive(Debug, Clone, Default)]
struct IndexMap<K: Eq + std::hash::Hash, V>(pub std::collections::HashMap<K, V>);

#[allow(dead_code)]
impl<K, V> std::ops::Deref for IndexMap<K, V>
where
    K: Eq + std::hash::Hash,
{
    type Target = std::collections::HashMap<K, V>;
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

#[allow(dead_code)]
impl<'a, K, V, Q: ?Sized> std::ops::Index<&'a Q> for IndexMap<K, V>
where
    K: std::borrow::Borrow<Q> + Eq + std::hash::Hash,
    Q: Ord + Eq + std::hash::Hash,
{
    type Output = V;

    fn index(&self, index: &Q) -> &V {
        self.0.get(index).expect("no entry found for key")
    }
}

#[allow(dead_code)]
impl<'a, K, V, Q: ?Sized> std::ops::IndexMut<&'a Q> for IndexMap<K, V>
where
    K: std::borrow::Borrow<Q> + Eq + std::hash::Hash,
    Q: Ord + Eq + std::hash::Hash,
{
    fn index_mut(&mut self, index: &Q) -> &mut V {
        self.0.get_mut(index).expect("no entry found for key")
    }
}

上で定義したIndexMapはDerefによって中身のHashMapにすることもできる。 これで let mut a = map![0; 0..20; "abc".chars()]; とかすれば a[&0][&'b'] = 20 とかして配列のようにアクセスできるのでdpだって書けてしまう。

調子にのって大きめの定義域で回すと遅い

Submission #5311611 - AtCoder Beginner Contest 122

これは上とはmapマクロの定義が少し違うが、定義域がchar5個分しかないので普通にACできている。

Submission #5400783 - AtCoder Beginner Contest 082

でもこっちは -8000..8000 でdpを回して、TLEしてしまう。 Vec だと600msくらいでAC。調子に乗ってはいけない。