#848758 (Go) No.1332 Range Nearest Query

提出ソース

結果

問題	No.1332 Range Nearest Query
ユーザー	草苺奶昔
提出日時	2023-03-13 12:12:23
言語	Go (1.22.1)
結果	AC
実行時間	841 ms / 2,500 ms
コード長	7,861 bytes
コンパイル時間	10,004 ms
コンパイル使用メモリ	228,396 KB
実行使用メモリ	16,280 KB
最終ジャッジ日時	2024-09-18 07:26:13
合計ジャッジ時間	41,042 ms
ジャッジサーバーID （参考情報）	judge3 / judge5

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	1 ms 6,816 KB
testcase_01	AC	2 ms 6,944 KB
testcase_02	AC	2 ms 6,944 KB
testcase_03	AC	797 ms 16,244 KB
testcase_04	AC	746 ms 16,280 KB
testcase_05	AC	738 ms 16,172 KB
testcase_06	AC	648 ms 12,292 KB
testcase_07	AC	689 ms 12,296 KB
testcase_08	AC	664 ms 12,296 KB
testcase_09	AC	647 ms 12,292 KB
testcase_10	AC	653 ms 12,292 KB
testcase_11	AC	659 ms 12,292 KB
testcase_12	AC	648 ms 12,300 KB
testcase_13	AC	656 ms 12,296 KB
testcase_14	AC	648 ms 12,300 KB
testcase_15	AC	660 ms 12,292 KB
testcase_16	AC	802 ms 14,268 KB
testcase_17	AC	785 ms 14,216 KB
testcase_18	AC	778 ms 14,268 KB
testcase_19	AC	796 ms 12,300 KB
testcase_20	AC	782 ms 12,300 KB
testcase_21	AC	841 ms 12,280 KB
testcase_22	AC	838 ms 14,332 KB
testcase_23	AC	809 ms 14,268 KB
testcase_24	AC	808 ms 12,292 KB
testcase_25	AC	769 ms 14,332 KB
testcase_26	AC	418 ms 14,336 KB
testcase_27	AC	509 ms 14,216 KB
testcase_28	AC	313 ms 6,940 KB
testcase_29	AC	314 ms 6,940 KB
testcase_30	AC	318 ms 6,940 KB
testcase_31	AC	316 ms 6,944 KB
testcase_32	AC	315 ms 6,944 KB
testcase_33	AC	316 ms 6,940 KB
testcase_34	AC	303 ms 6,944 KB
testcase_35	AC	304 ms 6,944 KB
testcase_36	AC	307 ms 6,944 KB
testcase_37	AC	318 ms 6,940 KB
testcase_38	AC	652 ms 10,172 KB
testcase_39	AC	484 ms 7,836 KB
testcase_40	AC	771 ms 12,284 KB
testcase_41	AC	535 ms 8,064 KB
testcase_42	AC	636 ms 10,168 KB
testcase_43	AC	582 ms 7,908 KB
testcase_44	AC	726 ms 14,360 KB
testcase_45	AC	704 ms 14,256 KB
testcase_46	AC	607 ms 10,048 KB
testcase_47	AC	671 ms 10,180 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

package main

import (
	"bufio"
	"fmt"
	"math/bits"
	"os"
)

func main() {
	in := bufio.NewReader(os.Stdin)
	out := bufio.NewWriter(os.Stdout)
	defer out.Flush()

	var n int
	fmt.Fscan(in, &n)
	nums := make([]int, n)
	for i := 0; i < n; i++ {
		fmt.Fscan(in, &nums[i])
	}

	M := NewWaveletMatrix(nums, 32) // 1e9
	var q int
	fmt.Fscan(in, &q)
	for i := 0; i < q; i++ {
		var left, right, x int
		fmt.Fscan(in, &left, &right, &x)
		left--
		res := INF
		lower := M.Floor(left, right, x) // 严格小于x的最大值
		if lower != -INF {
			res = min(res, abs(lower-x))
		}
		higher := M.Ceil(left, right, x) // 严格大于x的最小值
		if higher != INF {
			res = min(res, abs(higher-x))
		}

		fmt.Fprintln(out, res)
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func abs(a int) int {
	if a < 0 {
		return -a
	}
	return a
}

const INF int = 1e18

// 指定された配列から WaveletMatrix を構築する.
//  data:変換する配列(data[i]は0以上、2^maxLog未満)
//  maxLog:queryの最大値のbit数(普通は32)
func NewWaveletMatrix(data []int, maxLog int) *WaveletMatrix {
	dataCopy := append(data[:0:0], data...)
	n := len(dataCopy)
	mat := make([]*BitVector, maxLog)
	zs := make([]int, maxLog)
	buff1 := make([]int, maxLog)
	buff2 := make([]int, maxLog)

	ls, rs := make([]int, n), make([]int, n)
	for dep := 0; dep < maxLog; dep++ {
		mat[dep] = NewBitVector(n + 1)
		p, q := 0, 0
		for i := 0; i < n; i++ {
			k := (dataCopy[i] >> uint(maxLog-dep-1)) & 1
			if k == 1 {
				rs[q] = dataCopy[i]
				mat[dep].Set(i)
				q++
			} else {
				ls[p] = dataCopy[i]
				p++
			}
		}

		zs[dep] = p
		mat[dep].Build()
		ls = dataCopy
		for i := 0; i < q; i++ {
			dataCopy[p+i] = rs[i]
		}
	}

	return &WaveletMatrix{
		n:      n,
		maxLog: maxLog,
		mat:    mat,
		zs:     zs,
		buff1:  buff1,
		buff2:  buff2,
	}
}

type WaveletMatrix struct {
	n            int
	maxLog       int
	mat          []*BitVector
	zs           []int
	buff1, buff2 []int
}

// [start, end) に含まれる value の個数を求める.
//  alias: Rank
func (w *WaveletMatrix) Count(start, end, value int) int {
	return w.count(value, end) - w.count(value, start)
}

// [start, end) に含まれる [value, upper) の個数を求める.
// alias: RangeFreq
func (w *WaveletMatrix) CountRange(start, end, lower, upper int) int {
	return w.freqDfs(0, start, end, 0, lower, upper)
}

// k(0-indexed) 番目の value の位置を求める.
//  alias: Select
func (w *WaveletMatrix) Index(value, k int) int {
	w.count(value, w.n)
	for dep := w.maxLog - 1; dep >= 0; dep-- {
		bit := (value >> uint(w.maxLog-dep-1)) & 1
		k = w.mat[dep].IndexWithStart(bit, k, w.buff1[dep])
		if k < 0 || k >= w.buff2[dep] {
			return -1
		}
		k -= w.buff1[dep]
	}
	return k
}

func (w *WaveletMatrix) IndexWithStart(value, k, start int) int {
	return w.Index(value, k+w.count(value, start))
}

// [start, end) に含まれる要素の中で k(0-indexed) 番目に大きいものを求める.
//  alias: Quantile
func (w *WaveletMatrix) KthMax(start, end, k int) int {
	if k < 0 || k >= end-start {
		return -1
	}
	res := 0
	for dep := 0; dep < w.maxLog; dep++ {
		p, q := w.mat[dep].Count(1, start), w.mat[dep].Count(1, end)
		if k < q-p {
			start = w.zs[dep] + p
			end = w.zs[dep] + q
			res |= 1 << uint(w.maxLog-dep-1)
		} else {
			k -= q - p
			start -= p
			end -= q
		}
	}
	return res
}

// [start, end) に含まれる要素の中で k(0-indexed) 番目に小さいものを求める.
//  alias: Rquantile
func (w *WaveletMatrix) KthMin(start, end, k int) int {
	return w.KthMax(start, end, end-start-k-1)
}

// [start, end) に含まれる要素の中で value の次に小さいものを求める.存在しない場合は -INF を返す.
//  value >= 0
//  alias: Pred
func (w *WaveletMatrix) Lower(start, end, value int) int {
	k := w.lt(start, end, value)
	if k != 0 {
		return w.KthMin(start, end, k-1)
	}
	return -INF
}

// [start, end) に含まれる要素の中で value より大きいものを求める.存在しない場合は INF を返す.
//  value >= 0
//  alias: Succ
func (w *WaveletMatrix) Higher(start, end, value int) int {
	k := w.le(start, end, value)
	if k == end-start {
		return INF
	}
	return w.KthMin(start, end, k)
}

// [start, end) に含まれる要素の中で value 以下のものを求める.存在しない場合は -INF を返す.
func (w *WaveletMatrix) Floor(start, end, value int) int {
	count := w.Count(start, end, value)
	if count > 0 {
		return value
	}
	return w.Lower(start, end, value)
}

// [start, end) に含まれる要素の中で value 以上のものを求める.存在しない場合は INF を返す.
func (w *WaveletMatrix) Ceil(start, end, value int) int {
	count := w.Count(start, end, value)
	if count > 0 {
		return value
	}
	return w.Higher(start, end, value)
}

func (w *WaveletMatrix) access(k int) int {
	res := 0
	for dep := 0; dep < w.maxLog; dep++ {
		bit := w.mat[dep].Get(k)
		res = (res << 1) | bit
		k = w.mat[dep].Count(bit, k) + w.zs[dep]*dep
	}
	return res
}

func (w *WaveletMatrix) count(value, end int) int {
	left, right := 0, end
	for dep := 0; dep < w.maxLog; dep++ {
		w.buff1[dep] = left
		w.buff2[dep] = right
		bit := (value >> uint(w.maxLog-dep-1)) & 1
		left = w.mat[dep].Count(bit, left) + w.zs[dep]*bit
		right = w.mat[dep].Count(bit, right) + w.zs[dep]*bit
	}
	return right - left
}

func (w *WaveletMatrix) freqDfs(d, left, right, val, a, b int) int {
	if left == right {
		return 0
	}
	if d == w.maxLog {
		if a <= val && val < b {
			return right - left
		}
		return 0
	}

	nv := (1 << uint(w.maxLog-d-1)) | val
	nnv := ((1 << uint(w.maxLog-d-1)) - 1) | nv
	if nnv < a || b <= val {
		return 0
	}
	if a <= val && nnv < b {
		return right - left
	}
	lc, rc := w.mat[d].Count(1, left), w.mat[d].Count(1, right)
	return w.freqDfs(d+1, left-lc, right-rc, val, a, b) + w.freqDfs(d+1, lc+w.zs[d], rc+w.zs[d], nv, a, b)
}

func (w *WaveletMatrix) ll(left, right, v int) (int, int) {
	res := 0
	for dep := 0; dep < w.maxLog; dep++ {
		w.buff1[dep] = left
		w.buff2[dep] = right
		bit := v >> uint(w.maxLog-dep-1) & 1
		if bit == 1 {
			res += right - left + w.mat[dep].Count(1, left) - w.mat[dep].Count(1, right)
		}
		left = w.mat[dep].Count(bit, left) + w.zs[dep]*bit
		right = w.mat[dep].Count(bit, right) + w.zs[dep]*bit
	}
	return res, right - left
}

func (w *WaveletMatrix) lt(left, right, v int) int {
	a, _ := w.ll(left, right, v)
	return a
}

func (w *WaveletMatrix) le(left, right, v int) int {
	a, b := w.ll(left, right, v)
	return a + b
}

type BitVector struct {
	n     int
	block []int
	sum   []int
}

func NewBitVector(n int) *BitVector {
	blockCount := (n + 63) >> 6
	return &BitVector{
		n:     n,
		block: make([]int, blockCount),
		sum:   make([]int, blockCount),
	}
}

func (f *BitVector) Set(i int) {
	f.block[i>>6] |= 1 << uint(i&63)
}

func (f *BitVector) Build() {
	for i := 0; i < len(f.block)-1; i++ {
		f.sum[i+1] = f.sum[i] + bits.OnesCount(uint(f.block[i]))
	}
}

func (f *BitVector) Get(i int) int {
	return (f.block[i>>6] >> uint(i&63)) & 1
}

// [0,end) に含まれる 1 の個数.
func (f *BitVector) Count(value, end int) int {
	if value == 1 {
		return f.count1(end)
	}
	return end - f.count1(end)
}

// [0,end) に k(0-indexed) 番目の value の位置を求める.
// 存在しない場合は -1 を返す.
func (f *BitVector) Index(value, k int) int {
	if k < 0 || f.Count(value, f.n) <= k {
		return -1
	}

	left, right := 0, f.n
	for right-left > 1 {
		mid := (left + right) >> 1
		if f.Count(value, mid) >= k+1 {
			right = mid
		} else {
			left = mid
		}
	}
	return right - 1
}

func (f *BitVector) IndexWithStart(value, k, start int) int {
	return f.Index(value, k+f.Count(value, start))
}

func (f *BitVector) count1(k int) int {
	mask := (1 << uint(k&63)) - 1
	return f.sum[k>>6] + bits.OnesCount(uint(f.block[k>>6]&mask))
}

yukicoder

結果

テストケース

ソースコード