#831880 (PyPy3) No.738 平らな農地

提出ソース

結果

問題	No.738 平らな農地
ユーザー	terasa
提出日時	2023-01-10 19:36:12
言語	PyPy3 (7.3.15)
結果	AC
実行時間	672 ms / 2,000 ms
コード長	10,960 bytes
コンパイル時間	172 ms
コンパイル使用メモリ	82,148 KB
実行使用メモリ	161,592 KB
最終ジャッジ日時	2024-06-01 07:55:31
合計ジャッジ時間	32,101 ms
ジャッジサーバーID （参考情報）	judge3 / judge1

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

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	71 ms 70,784 KB
testcase_01	AC	70 ms 71,616 KB
testcase_02	AC	70 ms 71,056 KB
testcase_03	AC	71 ms 70,380 KB
testcase_04	AC	70 ms 70,188 KB
testcase_05	AC	157 ms 81,308 KB
testcase_06	AC	165 ms 81,444 KB
testcase_07	AC	160 ms 81,520 KB
testcase_08	AC	146 ms 80,488 KB
testcase_09	AC	124 ms 79,388 KB
testcase_10	AC	106 ms 79,560 KB
testcase_11	AC	137 ms 80,468 KB
testcase_12	AC	125 ms 79,572 KB
testcase_13	AC	155 ms 81,268 KB
testcase_14	AC	130 ms 79,700 KB
testcase_15	AC	465 ms 148,924 KB
testcase_16	AC	478 ms 152,044 KB
testcase_17	AC	509 ms 149,600 KB
testcase_18	AC	502 ms 147,348 KB
testcase_19	AC	543 ms 155,456 KB
testcase_20	AC	472 ms 150,624 KB
testcase_21	AC	530 ms 155,716 KB
testcase_22	AC	487 ms 151,976 KB
testcase_23	AC	529 ms 154,064 KB
testcase_24	AC	522 ms 151,900 KB
testcase_25	AC	121 ms 79,772 KB
testcase_26	AC	130 ms 80,236 KB
testcase_27	AC	130 ms 80,008 KB
testcase_28	AC	125 ms 79,620 KB
testcase_29	AC	125 ms 79,700 KB
testcase_30	AC	124 ms 80,068 KB
testcase_31	AC	126 ms 79,768 KB
testcase_32	AC	123 ms 79,792 KB
testcase_33	AC	129 ms 79,812 KB
testcase_34	AC	126 ms 79,852 KB
testcase_35	AC	128 ms 79,760 KB
testcase_36	AC	128 ms 79,996 KB
testcase_37	AC	126 ms 79,804 KB
testcase_38	AC	125 ms 79,384 KB
testcase_39	AC	125 ms 79,484 KB
testcase_40	AC	122 ms 79,728 KB
testcase_41	AC	122 ms 79,568 KB
testcase_42	AC	128 ms 80,336 KB
testcase_43	AC	126 ms 79,956 KB
testcase_44	AC	133 ms 80,192 KB
testcase_45	AC	477 ms 158,492 KB
testcase_46	AC	464 ms 151,668 KB
testcase_47	AC	467 ms 158,344 KB
testcase_48	AC	435 ms 152,992 KB
testcase_49	AC	439 ms 152,900 KB
testcase_50	AC	445 ms 158,352 KB
testcase_51	AC	490 ms 157,184 KB
testcase_52	AC	461 ms 152,656 KB
testcase_53	AC	463 ms 157,848 KB
testcase_54	AC	491 ms 156,940 KB
testcase_55	AC	485 ms 156,972 KB
testcase_56	AC	491 ms 158,588 KB
testcase_57	AC	452 ms 147,548 KB
testcase_58	AC	466 ms 159,328 KB
testcase_59	AC	467 ms 158,204 KB
testcase_60	AC	465 ms 158,352 KB
testcase_61	AC	469 ms 157,908 KB
testcase_62	AC	439 ms 150,380 KB
testcase_63	AC	496 ms 158,356 KB
testcase_64	AC	471 ms 158,088 KB
testcase_65	AC	474 ms 149,476 KB
testcase_66	AC	508 ms 154,288 KB
testcase_67	AC	491 ms 153,560 KB
testcase_68	AC	486 ms 156,040 KB
testcase_69	AC	634 ms 159,908 KB
testcase_70	AC	562 ms 160,448 KB
testcase_71	AC	116 ms 96,908 KB
testcase_72	AC	119 ms 97,856 KB
testcase_73	AC	121 ms 100,016 KB
testcase_74	AC	123 ms 100,016 KB
testcase_75	AC	573 ms 153,844 KB
testcase_76	AC	524 ms 156,660 KB
testcase_77	AC	618 ms 150,404 KB
testcase_78	AC	672 ms 160,540 KB
testcase_79	AC	585 ms 161,592 KB
testcase_80	AC	502 ms 157,096 KB
testcase_81	AC	551 ms 157,472 KB
testcase_82	AC	476 ms 156,872 KB
testcase_83	AC	285 ms 109,528 KB
testcase_84	AC	224 ms 109,576 KB
testcase_85	AC	582 ms 158,024 KB
testcase_86	AC	549 ms 154,924 KB
testcase_87	AC	305 ms 155,452 KB
testcase_88	AC	288 ms 155,272 KB
testcase_89	AC	68 ms 70,632 KB
testcase_90	AC	70 ms 71,208 KB
testcase_91	AC	71 ms 71,044 KB

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

ソースコード

raw source code

from typing import List, Tuple, Callable, TypeVar, Optional
import sys
import itertools
import heapq
import bisect
import math
from collections import deque, defaultdict
from functools import lru_cache, cmp_to_key

input = sys.stdin.readline

if __file__ != 'prog.py':
    sys.setrecursionlimit(10 ** 6)


def readints(): return map(int, input().split())
def readlist(): return list(readints())
def readstr(): return input()[:-1]
def readlist1(): return list(map(lambda x: int(x) - 1, input().split()))


class BitVector:
    # reference: https://tiramister.net/blog/posts/bitvector/
    TABLE = bytes([
        0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4,
        1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
        1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
        2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
        1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
        2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
        2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
        3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
        1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
        2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
        2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
        3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
        2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
        3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
        3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
        4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8,
    ])

    def __init__(self, N):
        self.cnum = (N + 255) >> 8

        self.bit = bytearray(self.cnum << 5)
        self.chunk = [0] * (self.cnum + 1)
        self.blocks = bytearray(self.cnum << 5)

        self.built = False

    def set(self, pos):
        self.bit[pos >> 3] |= 1 << (pos & 7)

    def access(self, pos):
        return self.bit[pos >> 3] >> (pos & 7) & 1

    def popcount(self, num):
        return self.TABLE[num]

    def build(self):
        for i in range(self.cnum):
            k = i << 5
            for j in range(31):
                self.blocks[k + 1] = self.blocks[k] + self.popcount(self.bit[k])
                k += 1
            self.chunk[i + 1] = self.chunk[i] + self.blocks[k] + self.popcount(self.bit[k])
        self.built = True

    def rank(self, pos):
        assert self.built
        cpos, tmp = pos >> 8, pos & 255
        bpos, offset = tmp >> 3, tmp & 7

        i = cpos << 5 | bpos
        rest = self.bit[i] & ((1 << offset) - 1)
        return self.chunk[cpos] + self.blocks[i] + self.popcount(rest)

    def select(self, num):
        """return minimum i that satisfies rank(i) = num"""
        assert self.built
        if num == 0:
            return 0
        if self.rank(self.N) < num:
            return -1

        l = -1
        r = self.N
        while r - l > 1:
            c = (l + r) >> 1
            if self.rank(c) >= num:
                r = c
            else:
                l = c
        return r


class WaveletMatrix:
    # reference: https://miti-7.hatenablog.com/entry/2018/04/28/152259
    def __init__(self, A: List[int], weight: List[int] = None):
        self.nums = sorted(set(A))
        self.idx = {a: i for i, a in enumerate(self.nums)}
        self.A = [self.idx[a] for a in A]

        self.digit = (len(self.nums) - 1).bit_length()
        self.B = [None] * self.digit
        self.offset = [None] * self.digit

        self.weight = weight
        if self.weight:
            self.S = [[0 for _ in range(len(self.A) + 1)] for _ in range(self.digit + 1)]
            for i, a in enumerate(self.A):
                self.S[self.digit][i + 1] = self.S[self.digit][i] + self.weight[i]

        if self.weight:
            T = list(zip(self.A, self.weight))
            for k in range(self.digit)[::-1]:
                self.B[k] = BitVector(len(T) + 1)
                zeros = []
                ones = []
                for i, (a, w) in enumerate(T):
                    if a >> k & 1:
                        self.B[k].set(i)
                        ones.append((a, w))
                    else:
                        zeros.append((a, w))
                self.B[k].build()
                self.offset[k] = len(zeros)
                T = zeros + ones
                for i, (a, w) in enumerate(T):
                    self.S[k][i + 1] = self.S[k][i] + w
        else:
            T = self.A
            for k in range(self.digit)[::-1]:
                self.B[k] = BitVector(len(T) + 1)
                zeros = []
                ones = []
                for i, a in enumerate(T):
                    if a >> k & 1:
                        self.B[k].set(i)
                        ones.append(a)
                    else:
                        zeros.append(a)
                self.B[k].build()
                self.offset[k] = len(zeros)
                T = zeros + ones

    def access(self, i: int):
        """return i-th value"""
        ret = 0
        cur = i
        for k in range(self.digit)[::-1]:
            if self.B[k].access(cur):
                ret |= (1 << k)
                cur = self.B[k].rank(cur) + self.offset[k]
            else:
                cur -= self.B[k].rank(cur)
        return self.nums[ret]

    def rank(self, l: int, r: int, x: int):
        """return the number of x's in [l, r) range"""
        x = self.idx.get(x)
        if x is None:
            return 0
        for k in range(self.digit)[::-1]:
            if x >> k & 1:
                l = self.B[k].rank(l) + self.offset[k]
                r = self.B[k].rank(r) + self.offset[k]
            else:
                l -= self.B[k].rank(l)
                r -= self.B[k].rank(r)
        return r - l

    def quantile(self, l: int, r: int, n: int):
        """return n-th (0-indexed) smallest value in [l, r) range"""
        assert 0 <= n < r - l
        ret = 0
        for k in range(self.digit)[::-1]:
            rank_l = self.B[k].rank(l)
            rank_r = self.B[k].rank(r)
            ones = rank_r - rank_l
            zeros = r - l - ones
            if zeros <= n:
                ret |= 1 << k
                l = rank_l + self.offset[k]
                r = rank_r + self.offset[k]
                n -= zeros
            else:
                l -= rank_l
                r -= rank_r
        return self.nums[ret]

    def rquantile(self, l: int, r: int, n: int):
        """return n-th (0-indeed) largest value in [l, r) range"""
        return self.quantile(l, r, r - l - 1 - n)

    def range_freq(self, l: int, r: int, lower: int, upper: int):
        """return the number of values s.t. lower <= x < upper in [l, r) range"""
        if lower >= upper:
            return 0
        return self._range_freq_upper(l, r, upper) - self._range_freq_upper(l, r, lower)

    def prev_value(self, l: int, r: int, upper: int):
        """return maximum x s.t. x < upper in [l, r) range if exist, otherwise None"""
        cnt = self._range_freq_upper(l, r, upper)
        if cnt == 0:
            return None
        return self.quantile(l, r, cnt - 1)

    def next_value(self, l: int, r: int, lower: int):
        """return minimum x s.t. x >= lower in [l, r) range if exist, otherwise None"""
        cnt = self._range_freq_upper(l, r, lower)
        if cnt == r - l:
            return None
        return self.quantile(l, r, cnt)

    def range_sum(self, l: int, r: int, lower: int, upper: int):
        """return sum of values s.t. lower <= x < upper in [l, r) range
        must be constructed with weight
        """
        assert self.weight
        return self._range_sum_upper(l, r, upper) - self._range_sum_upper(l, r, lower)

    def range_nsmallest_sum(self, l: int, r: int, n: int):
        """return sum of the first n (**1-indexed**) values ordered by ascending order in [l, r) range
        must be constructed with weight
        """
        assert self.weight
        if l >= r:
            return 0
        if self.digit == 0:
            return self.nums[0] * n
        ret = 0
        for k in range(self.digit)[::-1]:
            rank_l = self.B[k].rank(l)
            rank_r = self.B[k].rank(r)
            ones = rank_r - rank_l
            zeros = r - l - ones
            if zeros <= n:
                ret += self._get_internal_sum(k, l - rank_l, r - rank_r)
                l = rank_l + self.offset[k]
                r = rank_r + self.offset[k]
                n -= zeros
                if n == 0:
                    return ret
            else:
                l -= rank_l
                r -= rank_r
        ret += self._get_internal_sum(0, l, l + n)
        return ret

    def range_nlargest_sum(self, l: int, r: int, n: int):
        """return sum of the first n (**1-indexed**) values ordered by descending order in [l, r) range
        must be constructed with weight
        """
        assert self.weight
        return self._get_internal_sum(self.digit, l, r) - self.range_nsmallest_sum(l, r, r - l - n)

    def _range_freq_upper(self, l: int, r: int, upper: int):
        """return the number of values s.t. x < upper in [l, r) range"""
        if l >= r:
            return 0
        if upper > self.nums[-1]:
            return r - l
        if upper <= self.nums[0]:
            return 0
        upper = bisect.bisect_left(self.nums, upper)
        ret = 0
        for k in range(self.digit)[::-1]:
            rank_l = self.B[k].rank(l)
            rank_r = self.B[k].rank(r)
            ones = rank_r - rank_l
            zeros = r - l - ones
            if upper >> k & 1:
                ret += zeros
                l = rank_l + self.offset[k]
                r = rank_r + self.offset[k]
            else:
                l -= rank_l
                r -= rank_r
        return ret

    def _range_sum_upper(self, l: int, r: int, upper: int):
        """return sum of values s.t. x < upper in [l, r) range"""
        if l >= r:
            return 0
        if upper > self.nums[-1]:
            return self._get_internal_sum(self.digit, l, r)
        if upper <= self.nums[0]:
            return 0
        upper = bisect.bisect_left(self.nums, upper)
        ret = 0
        for k in range(self.digit)[::-1]:
            rank_l = self.B[k].rank(l)
            rank_r = self.B[k].rank(r)
            ones = rank_r - rank_l
            zero = r - l - ones
            if upper >> k & 1:
                ret += self._get_internal_sum(k, l - rank_l,  r - rank_r)
                l = rank_l + self.offset[k]
                r = rank_r + self.offset[k]
            else:
                l -= rank_l
                r -= rank_r
        return ret

    def _get_internal_sum(self, k: int, l: int, r: int):
        return self.S[k][r] - self.S[k][l]


N, K = readints()
A = readlist()
W = WaveletMatrix(A, A)
ans = 1 << 50
for i in range(N - K + 1):
    med = W.quantile(i, i + K, K // 2)
    v = W.range_nlargest_sum(i, i + K, K // 2) - W.range_nsmallest_sum(i, i + K, K // 2)
    ans = min(ans, v)
print(ans)

yukicoder

結果

テストケース

ソースコード