[PAT-A] 1057 Stack

Java

测试点1、2、3超时

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

public class Main {
    static int n;
    static Stack<Integer> s = new Stack<>();
    static TreeMap<Integer, Integer> sortedMap = new TreeMap<>();

    public static void main(String[] args) throws IOException {
        BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));

        n = Integer.parseInt(reader.readLine().trim());

        for (int i = 0; i < n; i++) {
            String command = reader.readLine().trim();
            if (command.startsWith("Push")) {
                int x = Integer.parseInt(command.split(" ")[1]);
                s.push(x);
                sortedMap.put(x, sortedMap.getOrDefault(x, 0) + 1);
            } else if (command.equals("Pop")) {
                if (s.isEmpty()) {
                    System.out.println("Invalid");
                } else {
                    int x = s.pop();
                    System.out.println(x);
                    if (sortedMap.get(x) == 1) {
                        sortedMap.remove(x);
                    } else {
                        sortedMap.put(x, sortedMap.get(x) - 1);
                    }
                }
            } else if (command.equals("PeekMedian")) {
                if (s.isEmpty()) {
                    System.out.println("Invalid");
                } else {
                    int medianIndex = (s.size() + 1) / 2;
                    int cnt = 0;
                    for (Map.Entry<Integer, Integer> entry : sortedMap.entrySet()) {
                        cnt += entry.getValue();
                        if (cnt >= medianIndex) {
                            System.out.println(entry.getKey());
                            break;
                        }
                    }
                }
            }
        }
        reader.close();
    }
}

C++

#include <iostream>
#include <cstring>
#include <algorithm>
#include <set>
#include <stack>

using namespace std;

int n;
stack<int> s;
multiset<int> l, r;        // 左小右大，mid = l.LastElem, l.size > r.size

void adjust() {
    while (r.size() > l.size()) {
        l.insert(*r.begin());
        r.erase(r.begin());
    }

    while (l.size() > r.size() + 1) {
        auto it = l.end();
        it--;
        r.insert(*it);
        l.erase(it);
    }
}

int main() {
    scanf("%d", &n);

    string op;
    while (n--) {
        cin >> op;

        if (op == "Push") {
            int x;
            scanf("%d", &x);
            s.push(x);

            if (r.empty() || x < *r.begin()) l.insert(x);
            else r.insert(x);

            adjust();
        } else if (op == "Pop") {
            if (s.empty()) printf("Invalid\n");
            else {
                int x = s.top();
                s.pop();
                printf("%d\n", x);

                auto it = l.end();
                it--;
                if (x <= *it) l.erase(l.find(x));
                else r.erase(r.find(x));

                adjust();
            }
        } else {
            if (s.empty()) printf("Invalid\n");
            else {
                auto it = l.end();
                it--;
                printf("%d\n", *it);
            }
        }
    }

    return 0;
};