Leverage Copy

メモの公開場所

Goでセグメントツリーを(可能な限り)抽象化

セグメントツリーをもう少し取り回しが効くようにしたいなぁと思った*1ので、 他の方々のブログ等を参考にしながら書き直してみました*2

実装

詳細は後述しますが、tsutajさん・ei1333さん・beetさんのお三方の解説を大いに参考にさせていただきました。

変数名やメソッド名はそれぞれの記事から拝借しているため、若干キメラになっていて読みにくいかもしれません。

また、数学的な側面の理解が浅いので、モノイド・作用素モノイド周りの命名が特に気持ち悪いかもしれません*3

通常(遅延伝搬なし)

type T int // (T, f): Monoid

type SegmentTree struct {
  sz   int              // minimum power of 2
  data []T              // elements in T
  f    func(lv, rv T) T // T <> T -> T
  ti   T                // identity element of Monoid
}

func NewSegmentTree(
  n int, f func(lv, rv T) T, ti T,
) *SegmentTree {
  st := new(SegmentTree)
  st.ti = ti
  st.f = f

  st.sz = 1
  for st.sz < n {
    st.sz *= 2
  }

  st.data = make([]T, 2*st.sz-1)
  for i := 0; i < 2*st.sz-1; i++ {
    st.data[i] = st.ti
  }

  return st
}

func (st *SegmentTree) Set(k int, x T) {
  st.data[k+(st.sz-1)] = x
}

func (st *SegmentTree) Build() {
  for i := st.sz - 2; i >= 0; i-- {
    st.data[i] = st.f(st.data[2*i+1], st.data[2*i+2])
  }
}

func (st *SegmentTree) Update(k int, x T) {
  k += st.sz - 1
  st.data[k] = x

  for k > 0 {
    k = (k - 1) / 2
    st.data[k] = st.f(st.data[2*k+1], st.data[2*k+2])
  }
}

func (st *SegmentTree) Query(a, b int) T {
  return st.query(a, b, 0, 0, st.sz)
}

func (st *SegmentTree) query(a, b, k, l, r int) T {
  if r <= a || b <= l {
    return st.ti
  }

  if a <= l && r <= b {
    return st.data[k]
  }

  lv := st.query(a, b, 2*k+1, l, (l+r)/2)
  rv := st.query(a, b, 2*k+2, (l+r)/2, r)
  return st.f(lv, rv)
}

func (st *SegmentTree) Get(k int) T {
  return st.data[k+(st.sz-1)]
}

おそらくは、モノイドの定義( T, f, ti )を適切に書き換えることだけに注力すれば、うまく動くのではないかと思います。 AOJにある典型例(RMQ, RSQ)は検証済みです*4

例: yukicoder No.875 Range Mindex Query

問題のURL

yukicoderの解説にもある通り、 最小値に加えて、最小値が入っているインデックスをもたせた構造体を要素の型とすればOKです。

以下はコードの抜粋です(提出はこちら)。

type T struct {
    v   int
    idx int
}

func main() {
    n, q := ReadInt2()
    A := ReadIntSlice(n)

    f := func(lv, rv T) T {
        t := T{}
        if lv.v < rv.v {
            t.v = lv.v
            t.idx = lv.idx
        } else {
            t.v = rv.v
            t.idx = rv.idx
        }
        return t
    }
    ti := T{v: 1<<31 - 1, idx: -1}
    st := NewSegmentTree(n, f, ti)
    for i := 0; i < n; i++ {
        st.Set(i, T{v: A[i], idx: i})
    }
    st.Build()

    for i := 0; i < q; i++ {
        c, l, r := ReadInt3()
        if c == 1 {
            ol := st.Get(l - 1)
            or := st.Get(r - 1)
            ol.idx, or.idx = or.idx, ol.idx
            st.Update(l-1, or)
            st.Update(r-1, ol)
        } else {
            e := st.Query(l-1, r)
            fmt.Println(e.idx + 1)
        }
    }
}

遅延伝搬あり

// Assumption: T == E
type T int // (T, f): Monoid
type E int // (E, h): Operator Monoid

type LazySegmentTree struct {
  sz   int
  data []T
  lazy []E
  f    func(lv, rv T) T        // T <> T -> T
  g    func(to T, from E) T    // T <> E -> T (assignment operator)
  h    func(to, from E) E      // E <> E -> E (assignment operator)
  p    func(e E, length int) E // E <> N -> E
  ti   T
  ei   E
}

func NewLazySegmentTree(
  n int,
  f func(lv, rv T) T, g func(to T, from E) T,
  h func(to, from E) E, p func(e E, length int) E,
  ti T, ei E,
) *LazySegmentTree {
  lst := new(LazySegmentTree)
  lst.f, lst.g, lst.h, lst.p = f, g, h, p
  lst.ti, lst.ei = ti, ei

  lst.sz = 1
  for lst.sz < n {
    lst.sz *= 2
  }

  lst.data = make([]T, 2*lst.sz-1)
  lst.lazy = make([]E, 2*lst.sz-1)
  for i := 0; i < 2*lst.sz-1; i++ {
    lst.data[i] = lst.ti
    lst.lazy[i] = lst.ei
  }

  return lst
}

func (lst *LazySegmentTree) Set(k int, x T) {
  lst.data[k+(lst.sz-1)] = x
}

func (lst *LazySegmentTree) Build() {
  for i := lst.sz - 2; i >= 0; i-- {
    lst.data[i] = lst.f(lst.data[2*i+1], lst.data[2*i+2])
  }
}

func (lst *LazySegmentTree) propagate(k, length int) {
  if lst.lazy[k] != lst.ei {
    if k < lst.sz-1 {
      lst.lazy[2*k+1] = lst.h(lst.lazy[2*k+1], lst.lazy[k])
      lst.lazy[2*k+2] = lst.h(lst.lazy[2*k+2], lst.lazy[k])
    }
    lst.data[k] = lst.g(lst.data[k], lst.p(lst.lazy[k], length))
    lst.lazy[k] = lst.ei
  }
}

func (lst *LazySegmentTree) Update(a, b int, x E) T {
  return lst.update(a, b, x, 0, 0, lst.sz)
}

func (lst *LazySegmentTree) update(a, b int, x E, k, l, r int) T {
  lst.propagate(k, r-l)

  if r <= a || b <= l {
    return lst.data[k]
  }

  if a <= l && r <= b {
    lst.lazy[k] = lst.h(lst.lazy[k], x)
    lst.propagate(k, r-l)
    return lst.data[k]
  }

  lv := lst.update(a, b, x, 2*k+1, l, (l+r)/2)
  rv := lst.update(a, b, x, 2*k+2, (l+r)/2, r)
  lst.data[k] = lst.f(lv, rv)
  return lst.data[k]
}

func (lst *LazySegmentTree) Query(a, b int) T {
  return lst.query(a, b, 0, 0, lst.sz)
}

func (lst *LazySegmentTree) query(a, b, k, l, r int) T {
  lst.propagate(k, r-l)

  if r <= a || b <= l {
    return lst.ti
  }

  if a <= l && r <= b {
    return lst.data[k]
  }

  lv := lst.query(a, b, 2*k+1, l, (l+r)/2)
  rv := lst.query(a, b, 2*k+2, (l+r)/2, r)
  return lst.f(lv, rv)
}

func (lst *LazySegmentTree) Get(k int) T {
  return lst.Query(k, k+1)
}

大抵はそうだろうということで、混乱しないように Assumption: T == E みたいなことを書いてしまいました。 しかし、次の例に示すように、別にそうとも限らないですね。。

こちらもモノイド、作用素モノイド、 そしてそれらに関する関数オブジェクトを適切に決定することに注力しさえすればいいように書いたつもりです。

例: yukicoder No.876 Range Compress Query

問題のURL

yukicoderの解説での想定解法は「階差数列に着目して区間加算を2箇所の点加算で済むようにする」 というようなもの*5のようです。 しかし、区間加算を真に受けて遅延伝搬セグメントツリーでも解けます。

与えられた定義で計算される圧縮値の他、区間の両端の値をもたせた構造体を要素の型とすればOKです。 単位元は少し注意が必要です。

以下はコードの抜粋です(提出はこちら)。

const INF_BIT60 = 1 << 60

func main() {
    n, q := ReadInt2()
    A := ReadIntSlice(n)

    f := func(lv, rv T) T {
        t := T{}
        t.v = lv.v + rv.v
        if lv.r >= INF_BIT60 || rv.l >= INF_BIT60 {
        } else if lv.r != rv.l {
            t.v++
        }
        t.l, t.r = lv.l, rv.r
        return t
    }
    g := func(to T, from E) T {
        to.l += int(from)
        to.r += int(from)
        return to
    }
    h := func(to, from E) E {
        return to + from
    }
    p := func(e E, length int) E {
        return e
    }
    ti := T{v: 0, l: INF_BIT60, r: INF_BIT60}
    ei := 0
    lst := NewLazySegmentTree(n, f, g, h, p, ti, E(ei))

    for i := 0; i < n; i++ {
        lst.Set(i, T{v: 0, l: A[i], r: A[i]})
    }
    lst.Build()

    for i := 0; i < q; i++ {
        c := ReadInt()
        if c == 1 {
            l, r, x := ReadInt3()
            lst.Update(l-1, r, E(x))
        } else {
            l, r := ReadInt2()
            t := lst.Query(l-1, r)
            fmt.Println(t.v + 1)
        }
    }
}

type T struct {
    v, l, r int
}

type E int

参考

学習の高速道路が整備されていてありがたいなぁと感じました。

tsutajさんの解説

tsutajさんの解説記事(遅延伝搬なし)

tsutajさんの解説記事(遅延伝搬あり)

いずれもとてもわかりやすく、初学時にセグメントツリーの仕組みを理解するのにはかどりました*6

セグメント木の抽象化だけに絞るのであれば、tsutajさんのブログを理解すれば十分な気がします。

ei1333さん・beetさんの解説

ei1333さんの旧ブログ記事での解説

beetさんの旧記事での解説

遅延伝搬セグメントツリーの抽象化は独力ではできる気がしなかったため、非常に参考になりました。 お二方とも旧ブログの方を参照してしまって恐縮です。 特にbeetさんの記事の方は本当に丁寧に説明されていてわかりやすかったです。


少数の定義を渡すだけで正しく動いてくれて感動したので、個人的には満足しています。

まだ解いた問題が少なすぎるので「これでOK」とは言い難いですが、 とりあえずは運用して様子見してみようと思います。

*1:解説で「セグメントツリーでも解けます」とあるときにテンション下がるのは嫌ですよね。

*2:言語を変えただけで、基本的に写経です。

*3:アドリブでいじるときに、自分にとって意味を思い出しやすいように意識して書いたらこうなりました

*4:Tの型が複雑になるときのインスタンス化の速度への影響が、まだ十分に検証しきれていません。

*5:そちらの解法で解き直していないので間違っているかもしれません。

*6:蟻本だけでは不安だったところを補完するのに活用させていただきました

Codeforces Round No.602 Div.2 D2復習

コンテスト中に解けなかったものの復習です。

いろいろな解法(というよりも解くために用いるツールが多様)がありますが、 BITを使った方法が一番自分にとって与し易かったため、BITで解きました。

D2. Optimal Subsequences (Hard Version)

問題のURL

問題

n 要素からなる数列 A が与えられる。

この数列に対して部分列(間の要素を好きなだけ削除し、残ったものの順番を変えずに得られる数列)を考える。

ある部分列の長さ k が与えられたとき、以下の条件を満たす部分列はoptimalであるという。

  1. 考えられる長さ k のあらゆる部分列の中で、部分列の要素の総和が最大となる。
  2. 1の条件を満たす部分列の中で、元の数列の位置に関して辞書順最小である。

一方で、 m 個のクエリが与えられる。

各クエリは k, pos の2つの1以上の整数からなり、このクエリに対して、 「長さ k のoptimalな部分列における pos 番目の値」を答える必要がある。

すべてのクエリに対して、それらの順番どおりに答えよ。

制約:

  • 1 <= n, m <= 2*10^5
  • 1 <= A[i] <= 10^9
  • 1 <= k <= n, 1 <= pos <= k

解答

Easyバージョンは制約が小さいため、愚直な解法で通りますが、こちらは賢くクエリを処理する必要があります。

まず整理すると、 1の条件よりoptimalな部分列は、元の数列を降順にソートしその先頭 k 個の数を選択したものとなります。 ただし、2の条件より、同じ要素が元の数列に複数存在する場合、できるだけ元の数列において前の方の位置に登場したものを優先的に選ぶ必要があります。

よって、元の数列 A を、まずは要素の大きさを基準に降順となるようにし、値が同じ場合は元の数列における位置に関して昇順となるようにソートします。 各クエリに答えるためには、まずこのソートした列に対して、先頭の k 要素を取得します。 そして、取得した k 要素の中から、位置に関して pos 番目に小さい要素の値を出力すればOKです。

f:id:maguroguma:20200104015921j:plain
optimal subsequenceの求め方

ただし、これを愚直に行うためには、 k 要素取得するたびにそれらを位置に関して昇順ソートする必要があります。 これは許容できないため、工夫が必要です。

まず、クエリをすべて先読みし、 k が小さい順に答えていくという工夫ができます。 こうすることで、(許されるなら)ソートしたい集合が、単純に新しい要素がappendされていくだけとなり、シンプルになります。

しかしながら、それでも k が大きくなるたびにソートするわけには行かないため、 「要素の追加」と「特定の順番の値の取得」を高速に行う必要があります。

このような操作は「BIT上の二分探索」を活用することで可能なため(詳しくは後述)、各クエリを O(logn) で処理できます。

よって、全体で O(m*logn) で解くことができます。

※1500msecぐらいかかりました。

※Sorterがうざったいかもしれませんが、AtCoderと共通化したいため、このようにしています(こどふぉのGoはバージョンが新しいので実際はもっとスマートに書けるはずです。)

var n int
var A []int
var m int

type BinaryIndexedTree struct {
    bit     []int
    n       int
    minPow2 int
}

// n(>=1) is number of elements of original data
func NewBIT(n int) *BinaryIndexedTree {
    newBit := new(BinaryIndexedTree)

    newBit.bit = make([]int, n+1)
    newBit.n = n

    newBit.minPow2 = 1
    for {
        if (newBit.minPow2 << 1) > n {
            break
        }
        newBit.minPow2 <<= 1
    }

    return newBit
}

// Sum of [1, i](1-based)
func (b *BinaryIndexedTree) Sum(i int) int {
    s := 0

    for i > 0 {
        s += b.bit[i]
        i -= i & (-i)
    }

    return s
}

// Add x to i(1-based)
func (b *BinaryIndexedTree) Add(i, x int) {
    for i <= b.n {
        b.bit[i] += x
        i += i & (-i)
    }
}

// LowerBound returns minimum i such that bit.Sum(i) >= w.
func (b *BinaryIndexedTree) LowerBound(w int) int {
    if w <= 0 {
        return 0
    }

    x := 0
    for k := b.minPow2; k > 0; k /= 2 {
        if x+k <= b.n && b.bit[x+k] < w {
            w -= b.bit[x+k]
            x += k
        }
    }

    return x + 1
}

func main() {
    n = ReadInt()
    A = ReadIntSlice(n)
    m = ReadInt()

    // クエリを先読みしてkで昇順ソート
    L := make(QueryList, 0)
    for i := 0; i < m; i++ {
        k, pos := ReadInt2()
        L = append(L, &Query{id: i, k: k, pos: pos}) // idは0-basedで格納
    }
    sort.Stable(L)

    // Aを値で降順→位置で昇順にソート
    S := make(ElementList, 0)
    for i := 0; i < n; i++ {
        S = append(S, &Element{val: A[i], pos: i}) // posは0-basedで格納
    }
    sort.Stable(S)

    // BITをOrderedSetとして扱う
    bit := NewBIT(n)
    ck := 0
    // クエリの回答保持用
    answers := make([]int, m)
    // クエリをkの小さい順に処理していく
    for i := 0; i < m; i++ {
        q := L[i]
        // BITに格納した要素数がk個になるまで追加する
        for ; ck < q.k; ck++ {
            e := S[ck]
            bit.Add(e.pos+1, 1)
        }

        // BIT中のpos番目を回答する
        ans := bit.LowerBound(q.pos)
        answers[q.id] = A[ans-1]
    }

    // まとめて回答
    for i := 0; i < m; i++ {
        fmt.Println(answers[i])
    }
}

type Query struct {
    id, k, pos int
}
type QueryList []*Query

func (l QueryList) Len() int {
    return len(l)
}
func (l QueryList) Swap(i, j int) {
    l[i], l[j] = l[j], l[i]
}

func (l QueryList) Less(i, j int) bool {
    return l[i].k < l[j].k
}

type Element struct {
    val, pos int
}
type ElementList []*Element

func (l ElementList) Len() int {
    return len(l)
}
func (l ElementList) Swap(i, j int) {
    l[i], l[j] = l[j], l[i]
}

func (l ElementList) Less(i, j int) bool {
    if l[i].val > l[j].val {
        return true
    } else if l[i].val < l[j].val {
        return false
    } else {
        return l[i].pos < l[j].pos
    }
}

BITを使って k 番目の要素を取得するテクニック

今回行ったような操作は競技プログラミングにおいては典型テクニックのようで、 例えばけんちょんさんのQiitaの記事などでも 詳しく解説されています(他にも色々わかりやすくまとめている方がたくさんおられました)。

今回は、位置に関してBITで管理するため、要素数の制約から座標圧縮する必要はありませんでした。

BIT上の二分探索とは?

hosさんのBIT解説PDFにおいて最後の方で説明がなされています。

。。が、自分には最初何をやっているのかよくわかりませんでした。

自分なりに議論を補間しつつゆっくりと追っていくと理解できたので、忘れた頃の未来の自分を第一の対象として、理解の道筋を残しておこうと思います。

まず、「BIT上の二分探索ってそもそも何?」となってしまいましたが、これは「元の配列の先頭からの累積和に関する二分探索」となります。 別の表現をすると 「1-basedなインデックス i で、 [1, i] の要素の総和が w 以上となる最小の i を探索する」 といえます。

二分探索自体の計算量が対数オーダーで、BITを使って累積和を求めるのも対数オーダーであるため、 計算量が log の2乗になるというのは納得できます。

しかし、BITの木構造をうまく利用することで、この二分探索も全体で O(logn) に落とすことができる、とのことです。 それとともに与えられたアルゴリズムが以下のものですが、これまた初見時はよくわかりませんでした。

func (b *BinaryIndexedTree) LowerBound(w int) int {
    if w <= 0 {
        return 0
    }

  // b.minPow2は、n以下の最小の2べき
  x := 0
    for k := b.minPow2; k > 0; k /= 2 {
        if x+k <= b.n && b.bit[x+k] < w {
            w -= b.bit[x+k]
            x += k
        }
    }

    return x + 1
}

これは要約すると、「先頭からの累積和を、できるだけ長い区間から足して良いかを都度判断する」ということを行っています。

流れとしては、BITの上方の区間が長いノードから順番に見ており、 k というのは現在注目している区間の長さを意味しています。

区間は、担当する区間の和を持っており、それがkey値 w よりも小さい場合は、足してから右側の次の短い長さの区間を見る必要があります。 逆にkey値よりも大きい場合は、その区間和は足さずに左側の次の短い長さの区間を見る必要があります。

以下は、ある具体例におけるアルゴリズムの様子を図示したものです(あまりわかりやすくできませんでしたが。。)。

f:id:maguroguma:20200104020050j:plain
BIT上の二分探索アルゴリズムの様子

結局、木の高さ分下ることで探索が終了するため、計算量は O(logn) となります。


クエリの先読みと、 k 番目に小さい要素の取得、という2つの典型テクニックが要求される問題だったと思います。

引用した記事でも言及されている平衡二分探索木についても、なるべく早く実装してみたいと思います。

次に出会ったときはちゃんと倒したいところです。

Codeforces Round No.594 Div.2 C復習

以前コンテストに参加して解けなかったものの復習です。

公式Editorialがハイコンテクスト過ぎてよくわからなかったのと、 数え上げの方法の典型度合いがものすごく高い気がしたので、別記事として書きました。

C. Ivan the Fool and the Probability Theory

問題のURL

問題

nm 列からなるグリッドを、黒・白の2色で塗り分けることを考える。

ただし、塗り方は「辺に関して隣接しているグリッドについて、同じ色の隣接グリッドがたかだか1つまで」とする。

このような塗り方は何通り存在するか、 10^9 + 7 で割ったあまりで答えよ。

解答

一見複雑だが、「左上から少数のグリッドの色を確定させると、他のグリッドの色は自動的に決まる」というようなことはよくある(と思う)。

とりあえず、1行目を条件に反しないように適当に塗ってみると(同じ色が連続するのは2回までとする)、 以下の図のように、同じ色が連続してしまうと、その列を中心として次の行は一意に定まってしまう。

f:id:maguroguma:20200102215055j:plain
1行目で同じ色が連続する場合

よって、このようなバターンの総数を数え上げる必要がある(※)。

また、同じ色が連続せずに、黒・白が交互に塗られる場合、 以下の図のように、次の行の1列目の色を決めた時点で、他の列の色が確定してしまう。

f:id:maguroguma:20200102215123j:plain
1行目で異なる色が交互に塗られる場合

よって、このような場合も結局、1列目の色の塗り方の総数を数え上げる必要があり、 これは、(※)の部分と同じようにして数え上げることができる。

黒・白が2つまで連続して良いという条件のもとでの、1次元グリッドの塗り分けの総数を求めることに集中できる。

とりあえず、最初の色を黒として樹形図を書いてみる。 以下の図に示すように、よくよく観察すると、 色に関係なく2つ前と1つ前の通り数の和となるように遷移する(フィボナッチ数列)ため、 dp[i-2] + dp[i-1] = dp[i] のようにまとめられる。

f:id:maguroguma:20200102215212j:plain
DPによる数え上げ

以上をもとにコーディングしていく。

同じ色が隣接することがない塗り方1通りを2重に数えない、 左上を白とした場合も数えるために2倍する、あたりを忘れないように注意する。

var n, m int64
var dp [100000 + 5]int64

func main() {
    n, m = ReadInt64_2()

    dp[0], dp[1] = 1, 2
    for i := 2; i <= 100000+3; i++ {
        dp[i] = dp[i-1] + dp[i-2]
        dp[i] %= MOD
    }

    var ans int64
    ans = dp[m-1] - 1
    ans += dp[n-1]
    ans %= MOD
    ans *= 2
    ans %= MOD

    fmt.Println(ans)
}

上記のフィボナッチ数を考えるあたりって自明なんでしょうか? 自分にはあまり自明には感じられないので、説明がかなり冗長かもしれません。

Codeforces Round No.609 Div.2 参加記録 (A〜D解答)

青になれました!やったね!

f:id:maguroguma:20191222030152p:plain

思ったより長い戦いだった。。

Codeforcesの青はAtCoderの水色」ぐらいの言説を見た気がするので、 「まぁ青ぐらいなら楽勝だろう」と思ったのですが、想像の3倍ぐらいは難しかったです。

多分青を維持するのは今の自分にはきつくて、 Codeforcesの1600はAtCoderの1400(水色の真ん中)ぐらいなんだろうなぁと感じています。

とはいえ、ミスは減ってきて実装力も向上しているとは思うので、 今後もしばらくはこの調子で頑張りたいところです。

A. Equation

問題のURL

解答

式変形すると a = n + b であるため、小さい方である b が決まれば a も自然と決まる。 そして、 a, b がともに合成数になるようにする。

2 以外の偶数はすべて合成数であることを踏まえると、

  1. n が偶数ならば、 b を適当な 2 以上の偶数とすれば n + b も当然 2 以上の偶数となる。
  2. n が奇数ならば、 b を適当な合成数である奇数とすれば n + b2 以上の偶数となる。

であり、 a, b を選択できる。

var n int

func main() {
    n = ReadInt()

    if n%2 == 0 {
        b := 4
        a := n + b
        fmt.Println(a, b)
    } else {
        b := 9
        a := n + b
        fmt.Println(a, b)
    }
}

こんな問題に9分もかかってしまいました。

というより、公式Editorialを見ると、「9*n, 8*n でOK」とのこと、たしかに。

こどふぉのDiv2のA問題って、 O(1) 算数をしないといけなかったり、真面目に全探索しないといけなかったりで、 なかなか難しいです。

簡単と思える日が来てほしい。

B. Modulo Equality

問題のURL

解答

A, B それぞれについて、登場する値のカウントをメモしておく。

答えが x のとき、すべての A[i] に対して、 (AにおけるA[i]の個数) == (Bにおける(b = (A[i] + x) % m)の個数) となっている必要がある。

よって、このような x としてあり得るものをすべて列挙し、 それが条件を満たすか確かめれば良い。 ただし、求められている x は条件を満たす最小のものであるため、見つけるたびに小さい方で更新する。

x としてありうるものは、ある適当な A の値(以下のコード中では A[0] としている)を取り出し、 これが B のどの要素に遷移するのか?を考えることで列挙できる。

また、各 x に対して、すべての A の要素の遷移先の個数が、 A の遷移前の個数と一致するかをチェックすれば良い。

以上より、トータルの計算量は O(n^2) となり、十分高速である。

var n, m int
var A, B []int

func main() {
    n, m = ReadInt2()
    A = ReadIntSlice(n)
    B = ReadIntSlice(n)

    amemo, bmemo := make(map[int]int), make(map[int]int)
    for i := 0; i < n; i++ {
        amemo[A[i]]++
        bmemo[B[i]]++
    }

    ans := -1
    aval, anum := A[0], amemo[A[0]]
    for bval, bnum := range bmemo {
        if anum != bnum {
            continue
        }

        var x int
        if bval > aval {
            x = bval - aval
        } else {
            x = bval + m - aval
        }

        isOK := true
        for av, an := range amemo {
            bv := (av + x) % m
            bn := bmemo[bv]
            if an == bn {
                continue
            } else {
                isOK = false
                break
            }
        }

        if isOK {
            if ans == -1 {
                ans = x
            } else {
                ChMin(&ans, x)
            }
        }
    }

    if ans == m {
        ans = 0
    }

    fmt.Println(ans)
}

A[i], B[i] の範囲が 10^9 なので固定長配列ではなく辞書を使う必要があるわけですが、 あまり慣れていない操作だったので時間がかかってしまいました。

そろそろこどふぉも100問弱ぐらいACしているはずですが、まだまだ足りてないなぁと感じます。

C. Long Beautiful Integer

問題のURL

解答

自由が効くのは最初の k 桁までであり、以降の桁は k 桁までの設定に従属する。

できるだけ小さいBeautifulな数値を作ろうとすると、 貪欲に上の桁から A の桁をコピーする必要があるため、まずはそのように桁を埋めていく。

このようにして出来る数値が A 以上であるならば、それが答えとなる。

もし A 未満となってしまう場合は、自由が効く桁の一番下の桁、すなわち k 桁目から微調整することを考える。

k 桁目までは A のコピーになっているため、できるだけ下の方の桁で +1 できれば良い (必ず A よりも大きくなる)。 ただし、注目桁が 9 である場合は繰り上がりを考える必要がある。 その場合は、 9 の桁を 0 にした後、上の桁に移動してその桁の数字を +1 することを再考する。

1桁でも変更ができたらその時点で終了し、出来上がった数を出力すれば良い。

k 桁目までを変更する際は、下の方の桁も合わせて修正する必要があることには注意する。

var n, k int
var A []rune

func main() {
    n, k = ReadInt2()
    A = ReadRuneSlice()

    isBeau := true
    for i := 0; i < n; i++ {
        if i+k < n {
            if A[i] != A[i+k] {
                isBeau = false
                break
            }
        } else {
            break
        }
    }
    if isBeau {
        fmt.Println(len(A))
        fmt.Println(string(A))
        return
    }

    B := make([]rune, n)
    for i := 0; i < n; i++ {
        B[i] = 'x'
    }

    // iは起点
    for i := 0; i < k; i++ {
        digit := A[i]

        // jは飛び飛び
        for j := i; j < n; j += k {
            B[j] = digit
        }
    }

    isSmall := false
    for i := 0; i < n; i++ {
        if A[i] == B[i] {
            continue
        } else if A[i] < B[i] {
            break
        } else {
            isSmall = true
            break
        }
    }
    if !isSmall {
        fmt.Println(len(B))
        fmt.Println(string(B))
        return
    }

    // 大きくなるように正しく調整する
    for i := k - 1; i >= 0; i-- {
        if B[i] == '9' {
            // iからk飛びで0に更新
            for j := i; j < n; j += k {
                B[j] = '0'
            }
        } else {
            // iからk飛びで更新
            for j := i; j < n; j += k {
                B[j]++
            }
            break
        }
    }

    fmt.Println(len(B))
    fmt.Println(string(B))
}

桁DPでよくやる考え方の基本、という感じがしました。

多分、最初の部分の「A がすでにBeautifulかどうか?」の判定は要らないと思います。

桁数を出力するのを忘れて1WA、 isSmall の判定部分でバグっており2WAしてしまい、 もったいないことをしてしまいました。

実装がちょっと長くなってくるとこういったミスが増えてくるので、 なにか対策を考えます。

D. Domino for Young

問題のURL

解答

わからなかったため、公式Editorialの手法をなぞりました。

与えられたヤング図形市松模様(チェッカーフラグ)のように塗ります。

その時の黒と白の数を数えて、小さい方が答えです。

小さい方の色を仮に黒とすると、黒は隣接する白とで1つのドミノによって専有され、 また、すべての未使用の黒は未使用の白とペアにすることができるため、とのこと。

賢い(けど、上の説明はちょっと大雑把すぎるかもしれないです)。

var n int
var A []int64
var b, w int64

func main() {
    n = ReadInt()
    A = ReadInt64Slice(n)

    b, w = 0, 0
    for i := 0; i < n; i++ {
        l := A[i]
        if i%2 == 0 {
            // Bから塗り始める
            b += (l + (2 - 1)) / 2
            w += l / 2
        } else {
            // Wから塗り始める
            w += (l + (2 - 1)) / 2
            b += l / 2
        }
    }

    fmt.Println(Min(b, w))
}

青になれて嬉しい気持ちはありますが、目標はまだまだ上に持てているので、 引き続き頑張っていく所存です。

Educational Codeforces Round No.78 参加記録 (A〜C解答)

いつもの自分だったらBの算数で詰まって終了だったので、 成長を喜びたいところ。

A. Shuffle Hashing

問題のURL

解答

ハッシュ文字列からパスワード分の長さの分だけ文字列を切り出して調べる、というのを全探索すれば良い。

切り出した部分とパスワードが一致するかどうかの判定は、文字配列をソートして 文字列として一致するかどうかを調べるのが簡単(だと思う)。

var t int
var P, H []rune

func main() {
    t = ReadInt()

    for tc := 0; tc < t; tc++ {
        P, H = ReadRuneSlice(), ReadRuneSlice()

        solve()
    }
}

func solve() {
    if len(P) > len(H) {
        fmt.Println("NO")
        return
    }

    pL := RuneList{}
    for i := 0; i < len(P); i++ {
        pL = append(pL, P[i])
    }
    sort.Sort(pL)

    plen := len(P)
    for i := 0; i < len(H); i++ {
        if i+plen > len(H) {
            break
        }

        tmp := H[i : i+plen]
        hL := RuneList{}
        for j := 0; j < len(tmp); j++ {
            hL = append(hL, tmp[j])
        }
        sort.Sort(hL)

        if string(pL) == string(hL) {
            fmt.Println("YES")
            return
        }
    }

    fmt.Println("NO")
}

type RuneList []rune

func (a RuneList) Len() int           { return len(a) }
func (a RuneList) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a RuneList) Less(i, j int) bool { return a[i] < a[j] }

ハッシュからパスワードを切り出す全探索の部分がコーディングに時間かかりすぎてしまっていたので、 ちょっと良くないですね。 どうしよう。

B. A and B

問題のURL

解答

2つの数 a, b を等しくするためには、最終的に2数の差 diff = abs(a-b) を埋める必要がある。

必要な操作回数を k 回としたときに、2つの数に加える数値の総和は k*(k+1)/2 である。 よって、この総和を2つの数に分割して、その2つの数の差が diff とすることができるのかを考える。

総和は好きなように分割できる一方で、分割した2つの数の差の偶奇は1種類のみである。

例: sum == 6 のとき (0, 6), (1, 5), (2, 4), (3, 3) が分割の全パターンで、2つの数の差はすべて偶数となる。

以上から、 k を小さいところから全探索して、 diff%2 == sum%2 && diff <= sum を満たすものを見つければ良い。

総和の形から1テストケースあたりの計算量は O(Sqrt(Max(a, b))) なので間に合う。

var t int
var a, b int64

func main() {
    t = ReadInt()

    for tc := 0; tc < t; tc++ {
        a, b = ReadInt64_2()

        solve()
    }
}

func solve() {
    diff := AbsInt(a - b)

    if diff == 0 {
        fmt.Println(0)
        return
    }

    for k := int64(1); ; k++ {
        sum := k * (k + 1) / 2

        if diff%2 == sum%2 && diff <= sum {
            fmt.Println(k)
            return
        }
    }
}

総和は好きなように分割できる一方で、分割した2つの数の差の偶奇は1種類のみである。

これって自明ですかね? 私は紙に色々書いているうちに「それはそうだ」と納得できる、ぐらいなんですが。

C. Berry Jam

問題のURL

解答

以下の図に示すように、いちごとブルーベリーそれぞれを 1, -1 に変換する(RはいちごでBはブルーベリー)。

f:id:maguroguma:20191221150527j:plain
与えられた情報のエンコーディング

ここで、与えられた配列の真ん中で分割する。 扱いやすくするため、食べられる順番に従うように、左側については元の配列の内側から並ぶように、配列を構築する。

この2つの配列に対して、累積和を考える。 このようにすると、問題は「左の累積和から1つ、右の累積和から1つ数を選んで、全体の合計と等しくなるようにせよ」 と読み替えることができる。 また、答えは左のインデックス、右のインデックスの和が小さいほど望ましい。

このように考えると、それぞれの配列についてある累積和を取るインデックスは、その最小のもののみに興味があるといえる (同じ数の累積和というのは、食べ残ったRとBの数の差が同じであり、等価であると言えるから(DPっぽい考え方))。

以下の実装では、それぞれの累積和に対する最小のインデックスを、左右の配列分それぞれ辞書で用意した (負数もまとめて入れたかったため)。

あとは、あり得る和のパターンの全探索を、左配列基準および右配列基準で2スキャンして、ベストなものを見つける。

var t int
var n int
var A []int

func main() {
    t = ReadInt()

    for tc := 0; tc < t; tc++ {
        n = ReadInt()
        A = ReadIntSlice(2 * n)

        solve()
    }
}

func solve() {
    E := make([]int, len(A))
    for i := 0; i < len(A); i++ {
        if A[i] == 2 {
            E[i] = -1
        } else {
            E[i] = 1
        }
    }
    sum := Sum(E...)

    if sum == 0 {
        fmt.Println(0)
        return
    }

    if sum < 0 {
        for i := 0; i < len(E); i++ {
            E[i] = -E[i]
        }
    }
    sum = Sum(E...)

    R, L := []int{}, []int{}
    for i := n - 1; i >= 0; i-- {
        L = append(L, E[i])
    }
    for i := 1; i < n; i++ {
        L[i] += L[i-1]
    }
    for i := n; i < 2*n; i++ {
        R = append(R, E[i])
    }
    for i := 1; i < n; i++ {
        R[i] += R[i-1]
    }

    leftMemo, rightMemo := make(map[int]int), make(map[int]int)
    leftMemo[0], rightMemo[0] = -1, -1
    for i := 0; i < n; i++ {
        a := L[i]
        if _, ok := leftMemo[a]; ok {
            continue
        } else {
            leftMemo[a] = i
        }
    }
    for i := 0; i < n; i++ {
        a := R[i]
        if _, ok := rightMemo[a]; ok {
            continue
        } else {
            rightMemo[a] = i
        }
    }

    ans := INF_BIT30
    for lv := 0; lv <= sum; lv++ {
        rv := sum - lv
        lidx, lok := leftMemo[lv]
        ridx, rok := rightMemo[rv]
        if lok && rok {
            ChMin(&ans, lidx+ridx+2)
        }
    }
    for rv := 0; rv <= sum; rv++ {
        lv := sum - rv
        lidx, lok := leftMemo[lv]
        ridx, rok := rightMemo[rv]
        if lok && rok {
            ChMin(&ans, lidx+ridx+2)
        }
    }
    fmt.Println(ans)
}

コンテスト後に時間を使って自分で考えた方法ですが、良い解法なのかはわかりません。 正直変換手数も多くてなんかまどろっこしい気がします。 editorialが公開された上で、模範解答が面白ければ、そちらも追記するかもしれません。


今回はなんとかイーブンに持ち込めました。

なんとか粘って今年の残りあと3回のこどふぉで青に持ち込みたいところです。

Codeforces Round No.608 (Div.2) 参加記録 (A〜C解答)

※Dはちょっと難しそうなので手を付けるのは当分先になりそうですが、 Eがシンプルな見た目で面白そうだったため、近々追記するかもしれません。

A. Suits

問題のURL

解答

2つのセットでジャケットが共通しており、ほかは独立している。 セットに組み込まれるアイテムはすべて1つずつであるため、 ジャケットはできるだけ高いセットで先に使ってしまい、 作れなくなったら、他方の安い方をできる限り作れば良い。

var a, b, c, d, e, f int

func main() {
    a, b, c, d = ReadInt4()
    e, f = ReadInt2()

    bc := Min(b, c)

    ans := 0
    if e > f {
        // aの方を使う
        m := Min(a, d)
        ans += e * m
        d -= m
        if d > 0 {
            mm := Min(bc, d)
            ans += f * mm
        }
    } else {
        // bcの方を使う
        m := Min(bc, d)
        ans += f * m
        d -= m
        if d > 0 {
            mm := Min(a, d)
            ans += e * mm
        }
    }

    fmt.Println(ans)
}

これもDiv.2のAの中では大分簡単な方な気がします。

B. Blocks

問題のURL

解答

初期の個数が偶数のものを反転させれば良い。

なぜなら、2つのブロックを反転させるとき、反転させて個数をゼロにしたい方の色のブロックの個数は、

  1. 2個減る
  2. 変わらない(1個減って1個増える)

のいずれかであることから、依然として偶数のままとなる。

最終的に選んだ色を0個にしなければならないことを考えると、 初期の個数が偶数の方しか選ぶことが出来ない。

具体的な反転回数については、左から愚直に選んだ色を反転させるようにすればよく、 高々 n 回で十分である。

var n int
var S []rune

func main() {
    n = ReadInt()
    S = ReadRuneSlice()

    // 最初から全部いっしょならOK
    memo := make(map[rune]int)
    memo['B'] = 0
    memo['W'] = 0
    for i := 0; i < n; i++ {
        memo[S[i]]++
    }
    if memo['B'] == n || memo['W'] == n {
        fmt.Println(0)
        return
    }

    // BもWも奇数なら無理
    if memo['B']%2 == 1 && memo['W']%2 == 1 {
        fmt.Println(-1)
        return
    }

    answers := []int{}
    if memo['B']%2 == 0 {
        // BをWにする
        for i := 0; i < n-1; i++ {
            if S[i] == 'B' {
                answers = append(answers, i+1)
                S[i] = 'W'
                S[i+1] = rev(S[i+1])
            }
        }
    } else {
        // WをBにする
        for i := 0; i < n-1; i++ {
            if S[i] == 'W' {
                answers = append(answers, i+1)
                S[i] = 'B'
                S[i+1] = rev(S[i+1])
            }
        }
    }

    fmt.Println(len(answers))
    fmt.Println(PrintIntsLine(answers...))
}

func rev(r rune) rune {
    if r == 'B' {
        return 'W'
    } else {
        return 'B'
    }
}

結構似たような問題は解いた気がするのに、かなり時間がかかってしまいました。

精進不足。

C. Shawarma Tent

問題のURL

解答

よくよく考えると、学校の1マス左・右・下・上のいずれかにテントを配置するのが最適です。

それぞれに配置した場合の、買い食いを行う可能性のある生徒の数を数え、 最大となるものを選べばよいです。

var n int
var sx, sy int
var X, Y []int

func main() {
    n = ReadInt()
    sx, sy = ReadInt2()

    X, Y = make([]int, n), make([]int, n)
    for i := 0; i < n; i++ {
        x, y := ReadInt2()
        X[i], Y[i] = x, y
    }

    l, r, u, d := 0, 0, 0, 0
    for i := 0; i < n; i++ {
        x, y := X[i], Y[i]

        if x <= sx-1 {
            l++
        }
        if x >= sx+1 {
            r++
        }
        if y <= sy-1 {
            d++
        }
        if y >= sy+1 {
            u++
        }
    }

    maxi := Max(l, r, u, d)
    fmt.Println(maxi)
    if l == maxi {
        fmt.Println(sx-1, sy)
    } else if r == maxi {
        fmt.Println(sx+1, sy)
    } else if d == maxi {
        fmt.Println(sx, sy-1)
    } else {
        fmt.Println(sx, sy+1)
    }
}

コンテスト中に通した解法では、片方の軸を学校に合わせれば良いということしか気づけず、 素直に区間で一番重なりが多い部分を求めるべく、家の位置のソート・ランレングス圧縮・累積和とかいう、 闇の深いコードを書いてしまいました(なぜ通ったのか)。

問題文がいかついと複雑に考えてしまうような気がするので、もうちょっと冷静に問題を俯瞰したいところです。

Codeforces Round No.607 (Div.2) 参加記録 (A〜D解答)

A. Suffix Three

問題のURL

解答

よくよく見ると末尾の2文字だけを見れば判定できるので、そこだけを見れば良い。

var t int
var S []rune

func main() {
    t = ReadInt()

    for tc := 0; tc < t; tc++ {
        S = ReadRuneSlice()

        solve()
    }
}

func solve() {
    n := len(S)
    str := S[n-2:]
    if string(str) == "po" {
        fmt.Println("FILIPINO")
    } else if string(str) == "su" {
        fmt.Println("JAPANESE")
    } else {
        fmt.Println("KOREAN")
    }
}

ここまで簡単なAはDiv.3でも見たことがないので新鮮でした。

B. Azamon Web Services

問題のURL

解答

1回までの入れ替えを許可されている文字列 S について、辞書順最小のものを求める。 それが C よりも辞書順で小さければ、それを出力すれば良い。 それが C よりも辞書順で大きいのならば、impossibleである。

O(n^2) が許される制約なので、愚直に考える。

できるだけ小さい文字をできるだけ前方に配置するのが良いので、 探索開始位置を先頭から末尾に1ずつ移動させていき、チェック位置の文字よりも小さいものが、 末尾までに存在すればそれと入れ替えるようにする。

ただし、最小のものが複数ある場合は、できる限り後の物を選ぶ必要があることに注意する。

var t int
var S, C []rune

func main() {
    t = ReadInt()

    for tc := 0; tc < t; tc++ {
        S = ReadRuneSlice()
        C = ReadRuneSlice()

        solve()
    }
}

const impossible = "---"

func solve() {
    n := len(S)
    lidx, ridx := -1, -1
    for i := 0; i < n; i++ {
        midx, minc := 5005, rune(100)
        // できるだけ後方から最小の文字を見つける
        for j := n - 1; j >= i+1; j-- {
            if minc > S[j] {
                midx, minc = j, S[j]
            }
        }

        if minc < S[i] {
            lidx, ridx = i, midx
            break
        }
    }

    if lidx == -1 && ridx == -1 {
        if string(S) < string(C) {
            fmt.Println(string(S))
        } else {
            fmt.Println(impossible)
        }
        return
    }

    S[lidx], S[ridx] = S[ridx], S[lidx]

    if string(S) < string(C) {
        fmt.Println(string(S))
    } else {
        fmt.Println(impossible)
    }
}

最初、最小のものを前から選ぶようなコードを書いてしまい、pretest2でWAしてしまって頭を抱えてしまいました。

何気なく WAWA という文字列を考えてたら(某プロゲーマーが浮かんだのかWAに怯えたのかは不明) 間違いに気づけたので良かったですが、 偶然以外何者でもないので、もう少し簡単なものでもいいから文字列問題に取り組むべきだなぁと思いました。

C. Cut and Paste

問題のURL

解答

x の値が初期状態の S の長さ以下であるときは素直にシミュレーションができる。

なぜなら、カットのあとすぐ行われるペーストは、切り出したものを少なくとも1回はペーストすることになるため、 x 文字目までは連続部分文字列が変わらず、長さだけに集中すれば良いからである。

ここで、 x の制約が 10^6 までであることを踏まえると、 S の長さが x に達するまでは文字列連結が必要だが、それ以降は不要となる。

後は実装を頑張るだけである。

var t int
var x int64
var S []rune

func main() {
    t = ReadInt()

    for tc := 0; tc < t; tc++ {
        x = ReadInt64()
        S = ReadRuneSlice()

        solve()
    }
}

var ans int64
var clip int64

func solve() {
    ans = int64(len(S))
    clip = 0

    for i := int64(0); i < x; i++ {
        c := int64(S[i] - '0')
        cut := S[i+1:]

        clip = NegativeMod(ans-(i+1), MOD)
        clip %= MOD
        ans += (c - 1) * (clip)
        ans %= MOD

        // 長さが十分だったらそれ以上は連結しない
        if int64(len(S)) < x {
            // 最大c-1回分追加で連結する
        OUTER:
            for j := int64(0); j < c-1; j++ {
                for k := 0; k < len(cut); k++ {
                    S = append(S, cut[k])

                    // x以上になったら連結しない
                    if int64(len(S)) >= x {
                        break OUTER
                    }
                }
            }
        }
    }

    fmt.Println(ans)
}

// NegativeMod can calculate a right residual whether value is positive or negative.
func NegativeMod(val, m int64) int64 {
    res := val % m
    if res < 0 {
        res += m
    }
    return res
}

添字が読みづらくて題意を把握するのに時間がかかってしまいました。

最近活用し始めましたが、Goのラベル付きfor文は競技プログラミングでは結構便利ですね (実開発では闇雲に使うと怒られそうですが)。

D. Beingawesomeism

問題のURL

解答

問題文がえげつないが、要はサンプル等で示された動きで A で塗りつぶしたい、ということ。

以下の図のように、1マスでも A が存在すればせいぜい4回の動作で塗りつぶしが可能となることを踏まえる。

f:id:maguroguma:20191218014546j:plain
必要な移動回数のイメージの素

少ない回数ほど条件は厳しいと考えられるので、 0, 1, 2, 3, 4回それぞれを図示して考えていく。

f:id:maguroguma:20191218014621j:plain
主要な場合分け

まず、図示していないが、全てが P の場合はimpossibleである。 また、すべてが A の場合は当然0回である。

次に、1回で済む場合を考えると、これは端っこの列もしくは行がすべて A のときである。 反対側へ向かう1回の移動で塗りつぶしが可能となる。

次に、2回で済む場合は、四隅のいずれかが A である場合と、 ある行もしくはある列についてすべてが A の場合である。 図に示す矢印の方向に移動を行えば、塗りつぶしが可能であることがわかる。

最後に、四隅を除いた端っこの行もしくは列に A が1文字でも存在すれば、3回の移動で済む。

いずれにも当てはまらない場合は、はじめの図に示したような4回の移動が必要となる。

var t int
var r, c int
var S [][]rune

func main() {
    t = ReadInt()

    for tc := 0; tc < t; tc++ {
        r, c = ReadInt2()
        S = [][]rune{}
        for i := 0; i < r; i++ {
            row := ReadRuneSlice()
            S = append(S, row)
        }

        solve()
    }
}

func solve() {
    anum := 0
    for i := 0; i < r; i++ {
        for j := 0; j < c; j++ {
            if S[i][j] == 'A' {
                anum++
            }
        }
    }
    if anum == 0 {
        fmt.Println("MORTAL")
        return
    }
    if anum == r*c {
        fmt.Println(0)
        return
    }

    if subRow(0) || subRow(r-1) || subCol(0) || subCol(c-1) {
        fmt.Println(1)
        return
    }

    b := false
    for i := 1; i < r-1; i++ {
        b = b || subRow(i)
    }
    for i := 1; i < c-1; i++ {
        b = b || subCol(i)
    }
    if b {
        fmt.Println(2)
        return
    }
    if S[0][0] == 'A' || S[0][c-1] == 'A' || S[r-1][0] == 'A' || S[r-1][c-1] == 'A' {
        fmt.Println(2)
        return
    }

    b = false
    for i := 1; i < c-1; i++ {
        b = b || (S[0][i] == 'A')
        b = b || (S[r-1][i] == 'A')
    }
    for i := 1; i < r-1; i++ {
        b = b || (S[i][0] == 'A')
        b = b || (S[i][c-1] == 'A')
    }
    if b {
        fmt.Println(3)
        return
    }

    fmt.Println(4)
}

func subRow(rowId int) bool {
    for i := 0; i < c; i++ {
        if S[rowId][i] != 'A' {
            return false
        }
    }
    return true
}

func subCol(colId int) bool {
    for i := 0; i < r; i++ {
        if S[i][colId] != 'A' {
            return false
        }
    }
    return true
}

pretestが弱くてなかなかの地獄だったようですが、 コンテスト外とはいえ1発で通ったので良かったです(残業後だったので通らなかったらやはり地獄だった)。

Cを解き終わった時点で満身創痍だったのでコンテスト中は触れられませんでしたが、 やっておくべきでした。