概要

2022-05-28(土)に開催されましたNOMURA プログラミングコンテスト2022（AtCoder Beginner Contest 253）のC - Max - Min QueryをGo言語で解く方法を紹介します。

問題と解答のアプローチ

Go言語で、問題文に沿って実装してみます。
$x$について下記のような制約がありますので、
問題文の多重集合$S$を、mapで書くことでメモリ制限を超えないようにします。

0 \leq x \leq 10^9

下記の書いてみますと、クエリ3で$S$の最大値と最小値を求めるところを
工夫をしないと、実行時間制限をオーバーしてしまいます。

package main

import (
	"bufio"
	"fmt"
	"os"
	"strconv"
)

var sc = bufio.NewScanner(os.Stdin)

type query struct {
	t, x, c int
}

func solveByMap(n int, q []query) (ans []int) {
	s := make(map[int]int)
	for _, v := range q {
		switch v.t {
		case 1:
			s[v.x]++
		case 2:
			if s[v.x] > v.c {
				s[v.x] -= v.c
			} else {
				delete(s, v.x)
			}
		case 3:
			//ここでmap[int]intの中から
			//max、minを効率よく求める必要があります。。
		}
	}
	return ans
}

func main() {
	buf := make([]byte, 1024*1024)
	sc.Buffer(buf, bufio.MaxScanTokenSize)
	sc.Split(bufio.ScanWords)

	//入力、データ型を整理
	n := nextInt()
	var q []query
	for i := 0; i < n; i++ {
		t := nextInt()
		switch t {
		case 1:
			q = append(q, query{t, nextInt(), 0})
		case 2:
			q = append(q, query{t, nextInt(), nextInt()})
		case 3:
			q = append(q, query{t, 0, 0})
		}
	}

	//解答
	ans := solveByMap(n, q)

	//出力
	out := bufio.NewWriter(os.Stdout)
	defer out.Flush()
	for _, v := range ans {
		fmt.Fprintln(out, v)
	}
}

func nextInt() int {
	sc.Scan()
	i, _ := strconv.Atoi(sc.Text())
	return i
}

C++では標準ライブラリにmultisetというデータ構造があり、
各クエリの操作を制限時間内で行える計算量で処理することができます。
しかし、Go言語の標準ライブラリに同じようなデータ構造が見受けられませんでした。

その関係か、C、D問題が同じくらいのDifficultyに対し、Go言語でのAC数はCの方が少ない(E問題と同じくらいの数)といった傾向がありました。

問題	Difficulty¹	競技時間中Go言語でAC数
A	17	37
B	57	32
C	518	13
D	520	22
E	1073	12

AtCoderのジャッジではRedBlackTreeのライブラリがインストールされています

私も過去の類題で知りました。
ABC170 E - Smart Infants

emirpasicさんのgods(Go Data Structures)がジャッジの環境にあります。
Go言語でAtCoderのコンテストに挑むには、標準以外自分で用意しており、
githubのライブラリが使えるというイメージがなかったため、初めて気づいた時には驚きました。

どのようなデータ構造になっているかは、リンク先の木構造が詳しく解説してくれています。
値の検索、追加、削除、再構築も概ね$O(log n)$の計算量で処理されます。
最大値は木の一番右側、最小値は木の一番左側で検索が可能です。

それを使って、下記のようなコードでACすることができました。

package main

import (
	"bufio"
	"fmt"
	"os"
	"strconv"

	"github.com/emirpasic/gods/trees/redblacktree"
)

var sc = bufio.NewScanner(os.Stdin)

type query struct {
	t, x, c int
}

func solve(n int, q []query) (ans []int) {
	s := redblacktree.NewWithIntComparator()

	insert := func(x int) {
		v, found := s.Get(x)
		if found {
			s.Put(x, v.(int)+1)
		} else {
			s.Put(x, 1)
		}
	}
	erase := func(x, c int) {
		v, found := s.Get(x)
		if found {
			if v.(int) > c {
				s.Put(x, v.(int)-c)
			} else {
				//v <= c
				s.Remove(x)
			}
		}
	}

	for _, v := range q {
		switch v.t {
		case 1:
			insert(v.x)
		case 2:
			erase(v.x, v.c)
		case 3:
			max := s.Right().Key.(int)
			min := s.Left().Key.(int)
			ans = append(ans, max-min)
		}
	}

	return ans
}

func main() {
	buf := make([]byte, 1024*1024)
	sc.Buffer(buf, bufio.MaxScanTokenSize)
	sc.Split(bufio.ScanWords)

	//入力、データ型を整理
	n := nextInt()
	var q []query
	for i := 0; i < n; i++ {
		t := nextInt()
		switch t {
		case 1:
			q = append(q, query{t, nextInt(), 0})
		case 2:
			q = append(q, query{t, nextInt(), nextInt()})
		case 3:
			q = append(q, query{t, 0, 0})
		}
	}

	//解答
	ans := solve(n, q)

	//出力
	out := bufio.NewWriter(os.Stdout)
	defer out.Flush()
	for _, v := range ans {
		fmt.Fprintln(out, v)
	}
}

func nextInt() int {
	sc.Scan()
	i, _ := strconv.Atoi(sc.Text())
	return i
}

残っている課題

以下の2問も同じようなアプローチで解こうとしました。

ABC241D - Sequence Queryの例がこちらです。
このコードは、記事を書いています2022-05-29現在、手元では動くけどもジャッジの環境ではコンパイルエラー(CE)となってしまいます。
これはgodsのredblacktree.IteratorAtが組み込まれたのが2020年頃でジャッジの環境はそれよりも前のバージョンと考えられます。
そのため、やはり自前でC++のmultisetのようなライブラリを準備しておく必要があるかもしれません。

package main

import (
	"bufio"
	"fmt"
	"os"
	"strconv"

	"github.com/emirpasic/gods/trees/redblacktree"
)

var sc = bufio.NewScanner(os.Stdin)

func Solve(q int, t, x, k []int) []int {

	IncrementTree := func(x int, t *redblacktree.Tree) {
		v, found := t.Get(x)
		if found {
			t.Put(x, v.(int)+1)
		} else {
			t.Put(x, 1)
		}
	}
	redBlackTree := redblacktree.NewWithIntComparator()
	var ans []int
	for i := 0; i < q; i++ {
		switch t[i] {
		case 1:
			IncrementTree(x[i], redBlackTree)
		case 2:
			node, _ := redBlackTree.Floor(x[i])
			it := redBlackTree.IteratorAt(node)
			s := it.Value().(int)
			found := s >= k[i]
			for !found && it.Prev() {
				s += it.Value().(int)
				found = s >= k[i]
			}
			if found {
				ans = append(ans, it.Key().(int))
			} else {
				ans = append(ans, -1)
			}
		case 3:
			node, _ := redBlackTree.Ceiling(x[i])
			it := redBlackTree.IteratorAt(node)
			s := it.Value().(int)
			found := s >= k[i]
			for !found && it.Next() {
				s += it.Value().(int)
				found = s >= k[i]
			}
			if found {
				ans = append(ans, it.Key().(int))
			} else {
				ans = append(ans, -1)
			}
		}
	}
	return ans
}

func main() {
	buf := make([]byte, 1024*1024)
	sc.Buffer(buf, bufio.MaxScanTokenSize)
	sc.Split(bufio.ScanWords)

	q := nextInt()
	var t, x, k []int
	for i := 0; i < q; i++ {
		t = append(t, nextInt())
		x = append(x, nextInt())
		if t[i] == 1 {
			k = append(k, -1)
		} else {
			k = append(k, nextInt())
		}
	}
	ans := Solve(q, t, x, k)
	out := bufio.NewWriter(os.Stdout)
	defer out.Flush()
	for _, v := range ans {
		fmt.Fprintln(out, v)
	}

}

func nextInt() int {
	sc.Scan()
	i, _ := strconv.Atoi(sc.Text())
	return i
}

参考

GoDS (Go Data Structures)

Atcoder Problemsより引用 ↩