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

提出ソース

結果

問題	No.738 平らな農地
ユーザー	terasa
提出日時	2023-01-09 13:20:29
言語	PyPy3 (7.3.15)
結果	AC
実行時間	1,170 ms / 2,000 ms
コード長	9,117 bytes
コンパイル時間	283 ms
コンパイル使用メモリ	87,024 KB
実行使用メモリ	175,504 KB
最終ジャッジ日時	2023-08-22 17:00:22
合計ジャッジ時間	54,958 ms
ジャッジサーバーID （参考情報）	judge11 / judge15

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

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	164 ms 79,948 KB
testcase_01	AC	162 ms 80,276 KB
testcase_02	AC	162 ms 80,080 KB
testcase_03	AC	163 ms 80,172 KB
testcase_04	AC	160 ms 79,948 KB
testcase_05	AC	308 ms 85,052 KB
testcase_06	AC	304 ms 86,072 KB
testcase_07	AC	280 ms 86,072 KB
testcase_08	AC	256 ms 84,764 KB
testcase_09	AC	230 ms 83,312 KB
testcase_10	AC	203 ms 82,660 KB
testcase_11	AC	247 ms 84,464 KB
testcase_12	AC	230 ms 83,224 KB
testcase_13	AC	290 ms 85,252 KB
testcase_14	AC	238 ms 84,060 KB
testcase_15	AC	715 ms 159,632 KB
testcase_16	AC	747 ms 161,608 KB
testcase_17	AC	817 ms 161,288 KB
testcase_18	AC	802 ms 159,432 KB
testcase_19	AC	947 ms 167,968 KB
testcase_20	AC	733 ms 161,788 KB
testcase_21	AC	874 ms 168,200 KB
testcase_22	AC	761 ms 162,384 KB
testcase_23	AC	863 ms 165,088 KB
testcase_24	AC	910 ms 162,508 KB
testcase_25	AC	237 ms 83,948 KB
testcase_26	AC	243 ms 84,060 KB
testcase_27	AC	242 ms 84,240 KB
testcase_28	AC	237 ms 83,380 KB
testcase_29	AC	246 ms 84,364 KB
testcase_30	AC	246 ms 84,092 KB
testcase_31	AC	230 ms 83,464 KB
testcase_32	AC	232 ms 84,144 KB
testcase_33	AC	239 ms 83,972 KB
testcase_34	AC	233 ms 83,932 KB
testcase_35	AC	228 ms 83,436 KB
testcase_36	AC	228 ms 83,500 KB
testcase_37	AC	232 ms 83,692 KB
testcase_38	AC	238 ms 83,892 KB
testcase_39	AC	233 ms 84,076 KB
testcase_40	AC	235 ms 84,044 KB
testcase_41	AC	239 ms 84,004 KB
testcase_42	AC	234 ms 83,752 KB
testcase_43	AC	230 ms 83,884 KB
testcase_44	AC	231 ms 83,684 KB
testcase_45	AC	713 ms 175,504 KB
testcase_46	AC	712 ms 160,144 KB
testcase_47	AC	723 ms 172,392 KB
testcase_48	AC	634 ms 163,356 KB
testcase_49	AC	648 ms 162,252 KB
testcase_50	AC	650 ms 172,200 KB
testcase_51	AC	736 ms 173,308 KB
testcase_52	AC	695 ms 162,832 KB
testcase_53	AC	696 ms 170,464 KB
testcase_54	AC	712 ms 173,356 KB
testcase_55	AC	743 ms 173,204 KB
testcase_56	AC	741 ms 171,740 KB
testcase_57	AC	695 ms 158,344 KB
testcase_58	AC	700 ms 171,516 KB
testcase_59	AC	687 ms 174,588 KB
testcase_60	AC	680 ms 175,316 KB
testcase_61	AC	690 ms 174,588 KB
testcase_62	AC	665 ms 160,916 KB
testcase_63	AC	750 ms 174,440 KB
testcase_64	AC	729 ms 174,844 KB
testcase_65	AC	1,053 ms 161,740 KB
testcase_66	AC	1,131 ms 166,208 KB
testcase_67	AC	764 ms 162,760 KB
testcase_68	AC	747 ms 165,596 KB
testcase_69	AC	993 ms 167,224 KB
testcase_70	AC	853 ms 169,064 KB
testcase_71	AC	206 ms 96,672 KB
testcase_72	AC	212 ms 96,680 KB
testcase_73	AC	209 ms 97,964 KB
testcase_74	AC	214 ms 98,300 KB
testcase_75	AC	931 ms 162,688 KB
testcase_76	AC	817 ms 165,820 KB
testcase_77	AC	996 ms 161,820 KB
testcase_78	AC	1,049 ms 168,816 KB
testcase_79	AC	1,013 ms 162,972 KB
testcase_80	AC	830 ms 165,952 KB
testcase_81	AC	925 ms 158,588 KB
testcase_82	AC	1,036 ms 161,264 KB
testcase_83	AC	379 ms 105,504 KB
testcase_84	AC	387 ms 104,188 KB
testcase_85	AC	1,170 ms 175,504 KB
testcase_86	AC	1,082 ms 161,188 KB
testcase_87	AC	351 ms 166,124 KB
testcase_88	AC	330 ms 162,392 KB
testcase_89	AC	162 ms 80,108 KB
testcase_90	AC	161 ms 80,004 KB
testcase_91	AC	161 ms 80,144 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):
        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.start_index = [-1] * len(self.nums)

        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.nums[a]

        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
            for i, a in enumerate(T):
                self.S[k][i + 1] = self.S[k][i] + self.nums[a]
        for i, a in enumerate(T):
            if self.start_index[a] < 0:
                self.start_index[a] = i

    def access(self, i):
        """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, i, x):
        """return the number of x's in [0, i) range"""
        x = self.idx.get(x)
        if x is None:
            return 0
        for k in range(self.digit)[::-1]:
            if x >> k & 1:
                i = self.B[k].rank(i) + self.offset[k]
            else:
                i -= self.B[k].rank(i)
        return i - self.start_index[x]

    def quantile(self, l, r, n):
        """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 range_freq(self, l, r, lower, upper):
        """return the number of values s.t. lower <= x < upper"""
        return self._range_freq_upper(l, r, upper) - self._range_freq_upper(l, r, lower)

    def prev_value(self, l, r, upper):
        """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, r, lower):
        """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, r, lower, upper):
        return self._range_sum_upper(l, r, upper) - self._range_sum_upper(l, r, lower)

    def range_sum_topn(self, l, r, n):
        assert 0 <= n < r - l
        if self.digit == 0:
            return self.nums[0] * (n + 1)
        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.S[k][r - rank_r] - self.S[k][l - rank_l]
                l = rank_l + self.offset[k]
                r = rank_r + self.offset[k]
                n -= zeros
            else:
                l -= rank_l
                r -= rank_r
        ret += self.S[0][l + n + 1] - self.S[0][l]
        return ret

    def _range_freq_upper(self, l, r, upper):
        """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, r, upper):
        if l >= r:
            return 0
        if upper > self.nums[-1]:
            return self.S[self.digit][r] - self.S[self.digit][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
            zero = r - l - ones
            if upper >> k & 1:
                ret += self.S[k][r - rank_r] - self.S[k][l - rank_l]
                l = rank_l + self.offset[k]
                r = rank_r + self.offset[k]
            else:
                l -= rank_l
                r -= rank_r
        return ret


N, K = readints()
A = readlist()
W = WaveletMatrix(A)
INF = 1 << 50
ans = INF
for i in range(N - K + 1):
    med = W.quantile(i, i + K, K // 2)
    upper_sum = W.range_sum(i, i + K, med, INF) - med * W.range_freq(i, i + K, med, INF)
    lower_sum = med * W.range_freq(i, i + K, 0, med) - W.range_sum(i, i + K, 0, med)
    ans = min(ans, upper_sum + lower_sum)
print(ans)

yukicoder

結果

テストケース

ソースコード