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

提出ソース

結果

問題	No.738 平らな農地
ユーザー	terasa
提出日時	2023-01-07 17:10:27
言語	PyPy3 (7.3.15)
結果	AC
実行時間	859 ms / 2,000 ms
コード長	7,376 bytes
コンパイル時間	251 ms
コンパイル使用メモリ	82,288 KB
実行使用メモリ	161,412 KB
最終ジャッジ日時	2024-05-08 12:38:32
合計ジャッジ時間	36,388 ms
ジャッジサーバーID （参考情報）	judge4 / judge5

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

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	67 ms 69,248 KB
testcase_01	AC	67 ms 69,120 KB
testcase_02	AC	66 ms 69,248 KB
testcase_03	AC	66 ms 69,504 KB
testcase_04	AC	67 ms 69,120 KB
testcase_05	AC	165 ms 81,740 KB
testcase_06	AC	169 ms 81,084 KB
testcase_07	AC	152 ms 81,544 KB
testcase_08	AC	127 ms 80,128 KB
testcase_09	AC	107 ms 79,708 KB
testcase_10	AC	92 ms 78,464 KB
testcase_11	AC	128 ms 79,872 KB
testcase_12	AC	109 ms 78,848 KB
testcase_13	AC	166 ms 81,608 KB
testcase_14	AC	111 ms 78,892 KB
testcase_15	AC	487 ms 147,968 KB
testcase_16	AC	514 ms 154,496 KB
testcase_17	AC	535 ms 149,376 KB
testcase_18	AC	526 ms 151,040 KB
testcase_19	AC	597 ms 159,232 KB
testcase_20	AC	498 ms 150,548 KB
testcase_21	AC	564 ms 159,104 KB
testcase_22	AC	531 ms 155,648 KB
testcase_23	AC	573 ms 157,440 KB
testcase_24	AC	579 ms 151,296 KB
testcase_25	AC	111 ms 79,104 KB
testcase_26	AC	112 ms 78,992 KB
testcase_27	AC	113 ms 79,616 KB
testcase_28	AC	110 ms 79,104 KB
testcase_29	AC	110 ms 79,392 KB
testcase_30	AC	109 ms 79,472 KB
testcase_31	AC	112 ms 78,992 KB
testcase_32	AC	109 ms 79,104 KB
testcase_33	AC	111 ms 79,104 KB
testcase_34	AC	110 ms 79,728 KB
testcase_35	AC	105 ms 78,876 KB
testcase_36	AC	113 ms 79,408 KB
testcase_37	AC	113 ms 78,976 KB
testcase_38	AC	108 ms 79,408 KB
testcase_39	AC	114 ms 79,744 KB
testcase_40	AC	114 ms 79,000 KB
testcase_41	AC	111 ms 79,356 KB
testcase_42	AC	109 ms 78,932 KB
testcase_43	AC	110 ms 78,988 KB
testcase_44	AC	109 ms 79,488 KB
testcase_45	AC	545 ms 157,056 KB
testcase_46	AC	515 ms 154,540 KB
testcase_47	AC	532 ms 154,496 KB
testcase_48	AC	500 ms 156,416 KB
testcase_49	AC	524 ms 156,032 KB
testcase_50	AC	517 ms 154,368 KB
testcase_51	AC	537 ms 155,776 KB
testcase_52	AC	519 ms 156,840 KB
testcase_53	AC	534 ms 153,984 KB
testcase_54	AC	524 ms 155,136 KB
testcase_55	AC	528 ms 155,264 KB
testcase_56	AC	516 ms 154,420 KB
testcase_57	AC	502 ms 150,784 KB
testcase_58	AC	501 ms 154,112 KB
testcase_59	AC	513 ms 157,312 KB
testcase_60	AC	489 ms 156,668 KB
testcase_61	AC	501 ms 156,672 KB
testcase_62	AC	488 ms 154,624 KB
testcase_63	AC	538 ms 157,252 KB
testcase_64	AC	516 ms 156,928 KB
testcase_65	AC	704 ms 149,468 KB
testcase_66	AC	744 ms 158,080 KB
testcase_67	AC	512 ms 152,376 KB
testcase_68	AC	510 ms 159,964 KB
testcase_69	AC	641 ms 160,812 KB
testcase_70	AC	568 ms 161,148 KB
testcase_71	AC	242 ms 103,856 KB
testcase_72	AC	234 ms 106,624 KB
testcase_73	AC	223 ms 106,624 KB
testcase_74	AC	254 ms 109,440 KB
testcase_75	AC	615 ms 151,380 KB
testcase_76	AC	567 ms 156,748 KB
testcase_77	AC	648 ms 154,136 KB
testcase_78	AC	698 ms 161,412 KB
testcase_79	AC	679 ms 154,952 KB
testcase_80	AC	589 ms 160,180 KB
testcase_81	AC	624 ms 153,296 KB
testcase_82	AC	620 ms 160,288 KB
testcase_83	AC	287 ms 101,120 KB
testcase_84	AC	262 ms 101,248 KB
testcase_85	AC	859 ms 157,896 KB
testcase_86	AC	741 ms 153,116 KB
testcase_87	AC	441 ms 155,088 KB
testcase_88	AC	418 ms 157,064 KB
testcase_89	AC	66 ms 69,376 KB
testcase_90	AC	65 ms 69,248 KB
testcase_91	AC	67 ms 68,992 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/
    # (reference implemention is succinct, but this implemention is not succinct.)
    def __init__(self, N):
        self.N = N
        self.block_num = (N + 31) >> 5
        self.bit = [0] * self.block_num
        self.sum = [0] * self.block_num

        self.built = False

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

    def build(self):
        assert not self.built
        for i in range(1, self.block_num):
            self.sum[i] = self.sum[i - 1] + self.popcount(self.bit[i - 1])
        self.built = True

    def access(self, pos):
        """return pos-th bit"""
        return self.bit[pos >> 5] >> (pos & 31) & 1

    def rank(self, pos):
        """count 1's in [0, pos) range"""
        assert self.built
        i = pos >> 5
        offset = pos & 31
        return self.sum[i] + self.popcount(self.bit[i] & ((1 << offset) - 1))

    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

    def popcount(self, n):
        c = (n & 0x5555555555555555) + ((n >> 1) & 0x5555555555555555)
        c = (c & 0x3333333333333333) + ((c >> 2) & 0x3333333333333333)
        c = (c & 0x0f0f0f0f0f0f0f0f) + ((c >> 4) & 0x0f0f0f0f0f0f0f0f)
        c = (c & 0x00ff00ff00ff00ff) + ((c >> 8) & 0x00ff00ff00ff00ff)
        c = (c & 0x0000ffff0000ffff) + ((c >> 16) & 0x0000ffff0000ffff)
        return c


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 = {}

        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.start_index.setdefault(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_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


class BIT:
    def __init__(self, N):
        self.N = N
        self.A = [0] * (N + 1)

    def add(self, i, x):
        """add x to i-th element (0-indexed)"""
        assert 0 <= i < self.N
        i += 1
        while i <= self.N:
            self.A[i] += x
            i += i & -i

    def sum(self, i):
        """return sum(A[:i])"""
        assert 0 <= i <= self.N
        s = 0
        while i > 0:
            s += self.A[i]
            i -= i & -i
        return s

    def range_sum(self, l, r):
        """return sum(A[l:r])"""
        return self.sum(r) - self.sum(l)


N, K = readints()
A = readlist()
WM = WaveletMatrix(A)

cands = [0] * (N - K + 1)
task = defaultdict(list)
for i in range(N - K + 1):
    med = WM.quantile(i, i + K, K // 2)
    task[med].append(i)

nums = WM.nums
idx = defaultdict(list)
for i, a in enumerate(A):
    idx[a].append(i)

bit = BIT(N)
bit2 = BIT(N)
for k in nums:
    for i in task[k]:
        cands[i] += k * bit2.range_sum(i, i + K) - bit.range_sum(i, i + K)
    for i in idx[k]:
        bit.add(i, k)
        bit2.add(i, 1)
bit = BIT(N)
bit2 = BIT(N)
for k in nums[::-1]:
    for i in task[k]:
        cands[i] += bit.range_sum(i, i + K) - k * bit2.range_sum(i, i + K)
    for i in idx[k]:
        bit.add(i, k)
        bit2.add(i, 1)
print(min(cands))

yukicoder

結果

テストケース

ソースコード