#1051895 (Java) No.738 平らな農地

• 問題: No.738 平らな農地
• ユーザー: norioc
• 提出日時: 2025-03-12 04:14:54
• 言語: Java (openjdk 23)
• 結果: AC
• 実行時間: 850 ms / 2,000 ms
• コード長: 7,700 bytes
• コンパイル時間: 2,651 ms
• コンパイル使用メモリ: 80,792 KB
• 実行使用メモリ: 65,888 KB
• 最終ジャッジ日時: 2025-03-12 04:15:44
• 合計ジャッジ時間: 45,590 ms
• ジャッジサーバーID （参考情報）: judge5 / judge1

ファイルパターン	結果
sample	AC * 5
other	AC * 87

ソースコード

import java.io.*;
import java.util.*;

public class Main {
    // ランダム化二分探索木（RBST）の実装
    static class RBST {
        // ノードクラス
        static class Node {
            int val;
            Node left, right;
            int size;
            long sum;

            Node(int val) {
                this.val = val;
                this.size = 1;
                this.sum = val;
                this.left = null;
                this.right = null;
            }

            // 今回は特に何もしない（遅延伝搬などが必要ならここに処理を追加）
            void update() {
                // no additional update needed
            }

            // 遅延伝搬用（必要なら実装）
            void push() {
                // no lazy propagation implemented
            }
        }

        Node root;
        static Random rand = new Random();

        public RBST() {
            root = null;
        }

        public RBST(Node node) {
            root = node;
        }

        // 木のサイズ（要素数）を返す
        public int size() {
            return _size(root);
        }

        // 木全体の合計値を返す
        public long sum() {
            return _sum(root);
        }

        private static long _sum(Node node) {
            return node == null ? 0L : node.sum;
        }

        private static int _size(Node node) {
            return node == null ? 0 : node.size;
        }

        private static Node _update(Node node) {
            if (node == null) return null;
            node.size = _size(node.left) + _size(node.right) + 1;
            node.sum = _sum(node.left) + _sum(node.right) + node.val;
            node.update();
            return node;
        }

        private static void _push(Node node) {
            if (node != null) {
                node.push();
            }
        }

        // 指定した値 "val" より小さい要素数を返す（lower_bound）
        private static int _lower_bound(Node node, int val) {
            if (node == null) return 0;
            _push(node);
            if (val <= node.val) {
                return _lower_bound(node.left, val);
            } else {
                return _size(node.left) + _lower_bound(node.right, val) + 1;
            }
        }

        public int lowerBound(int val) {
            return _lower_bound(root, val);
        }

        // 指定した値 "val" より大きい要素数（upper_bound）
        private static int _upper_bound(Node node, int val) {
            if (node == null) return 0;
            _push(node);
            if (val >= node.val) {
                return _size(node.left) + _upper_bound(node.right, val) + 1;
            } else {
                return _upper_bound(node.left, val);
            }
        }

        public int upperBound(int val) {
            return _upper_bound(root, val);
        }

        // 値 val の個数
        public int count(int val) {
            return upperBound(val) - lowerBound(val);
        }

        // ソート順で k 番目（0-indexed）の値を返す
        private static int _get(Node node, int k) {
            if (node == null) return -1; // エラー扱い（要改善）
            _push(node);
            int leftSize = _size(node.left);
            if (k == leftSize) {
                return node.val;
            } else if (k < leftSize) {
                return _get(node.left, k);
            } else {
                return _get(node.right, k - leftSize - 1);
            }
        }

        public int get(int k) {
            return _get(root, k);
        }

        // 2つの部分木をマージする
        private static Node _merge(Node left, Node right) {
            _push(left);
            _push(right);
            if (left == null || right == null) {
                return left != null ? left : right;
            }
            int totalSize = left.size + right.size;
            // 1～0xfffffff の乱数を生成
            int r = rand.nextInt(0xfffffff) + 1;
            if (r % totalSize < left.size) {
                left.right = _merge(left.right, right);
                return _update(left);
            } else {
                right.left = _merge(left, right.left);
                return _update(right);
            }
        }

        // 別のRBSTとマージする
        public void merge(RBST other) {
            root = _merge(root, other.root);
        }

        // _split の戻り値を保持するための Pair クラス
        static class Pair {
            Node first, second;
            Pair(Node first, Node second) {
                this.first = first;
                this.second = second;
            }
        }

        // ノードを位置 k（0-indexed）で分割する
        // [0, k) と [k, n) に分ける
        private static Pair _split(Node node, int k) {
            if (node == null) return new Pair(null, null);
            _push(node);
            if (k <= _size(node.left)) {
                Pair p = _split(node.left, k);
                node.left = p.second;
                return new Pair(p.first, _update(node));
            } else {
                Pair p = _split(node.right, k - _size(node.left) - 1);
                node.right = p.first;
                return new Pair(_update(node), p.second);
            }
        }

        // k 番目で分割し、右側を新たな RBST として返す
        public RBST split(int k) {
            Pair p = _split(root, k);
            root = p.first;
            return new RBST(p.second);
        }

        // 値 val を挿入する
        public void insert(int val) {
            int pos = lowerBound(val);
            Pair p = _split(root, pos);
            Node newNode = new Node(val);
            root = _merge(_merge(p.first, newNode), p.second);
        }

        // 値 val を1つ削除する
        public void remove(int val) {
            if (count(val) == 0) return;
            int pos = lowerBound(val);
            Pair p = _split(root, pos);
            Pair p2 = _split(p.second, 1);
            root = _merge(p.first, p2.second);
        }
    }

    static long INF = 1L << 60;

    // calcCost は現在の RBST 内の要素（窓の要素）に対してコストを計算する
    static long calcCost(RBST bst, int K) {
        if (K == 1) return 0;
        if (K == 2) {
            return (long) bst.get(1) - bst.get(0);
        }
        // 要素を半分に分割して中央値付近のコストを計算する
        RBST rt = bst.split(K / 2);
        int m = rt.get(0);

        long res = (long) m * bst.size() - bst.sum();
        res += (rt.sum() - m) - (long) m * (rt.size() - 1);
        bst.merge(rt);
        return res;
    }

    // main メソッド
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int N = Integer.parseInt(st.nextToken());
        int K = Integer.parseInt(st.nextToken());

        int[] A = new int[N];
        st = new StringTokenizer(br.readLine());
        for (int i = 0; i < N; i++) {
            A[i] = Integer.parseInt(st.nextToken());
        }

        long ans = INF;
        RBST bst = new RBST();
        for (int i = 0; i < N; i++) {
            bst.insert(A[i]);
            if (bst.size() > K) {
                bst.remove(A[i - K]);
            }
            if (bst.size() == K) {
                ans = Math.min(ans, calcCost(bst, K));
            }
        }

        System.out.println(ans);
    }
}