一道蚂蚁社招面试题：数据流的中位数

[Acceml]

前段时间面了蚂蚁一个技术还蛮强的团队,圈子比较小,所以在此不表。 二面面试官根据汇报层级推测应该是 P9 级别及以上，在美国电面我，面试风格偏硅谷那边。虽然感觉这道题面完自己没表现好，凉凉了，不过觉得还是蛮有意思，觉得自己思考问题还是不够严谨,有很大提高的空间，所以写出来总结。

题目

中位数是有序列表中间的数。如果列表长度是偶数，中位数则是中间两个数的平均值。 例如， [2,3,4] 的中位数是 3 [2,3] 的中位数是 (2 + 3) / 2 = 2.5

分析

这道题目是 Leetcode 的 hard 难度的原题。求中位数或者 topK 的问题我们通常都会想用堆来解决。 但是面试官又进一步加大了难度，他要求内存使用很小，没有磁盘，但是压榨空间的同时可以忍受一定时间的损耗。且面试官不仅要求说出思路，要写出完整可经过大数据检测的 production code。

先不考虑内存

不考虑内存的方式就是 Leetcode 论坛上的题解。 基本思想是建立两个堆。左边是大根堆，右边是小根堆。 如果是奇数的时候，大根堆的堆顶是中位数。

例如：[1,2,3,4,5] 大根堆建立如下：

      3
     / \
    1   2

小根堆建立如下:

      4
     / 
    5   

偶数的时候则是最大堆和最小堆顶的平均数。

例如: [1, 2, 3, 4]

大根堆建立如下：

      2
     / 
    1   

小根堆建立如下:

      3
     / 
    4   

然后再维护一个奇数偶数的状态即可求中位数。

public class MedianStream {
    private PriorityQueue<Integer> leftHeap = new PriorityQueue<>(5, Collections.reverseOrder());
    private PriorityQueue<Integer> rightHeap = new PriorityQueue<>(5);

    private boolean even = true;

    public double getMedian() {
        if (even) {
            return (leftHeap.peek() + rightHeap.peek()) / 2.0;
        } else {
            return leftHeap.peek();
        }
    }

    public void addNum(int num) {
        if (even) {
            rightHeap.offer(num);
            int rightMin = rightHeap.poll();
            leftHeap.offer(rightMin);
        } else {
            leftHeap.offer(num);
            int leftMax = leftHeap.poll();
            rightHeap.offer(leftMax);
        }
        System.out.println(leftHeap);
        System.out.println(rightHeap);
        // 奇偶变换.
        even = !even;
    }
}

压榨内存

但是这样做的问题就是可能内存会爆掉。如果你的流无限大，那么意味着这些数据都要存在内存中，堆必须要能够建无限大。如果内存必须很小的方式，用时间换空间。

• 流是可以重复去读的, 用时间换空间;
• 可以用分治的思想,先读一遍流,把流中的数据个数分桶;
• 分桶之后遍历桶就可以得到中位数落在哪个桶里面,这样就把问题的范围缩小了。

public class Median {
    private static int BUCKET_SIZE = 1000;

    private int left = 0;
    private int right = Integer.MAX_VALUE;

    // 流这里用 int[] 代替
    public double findMedian(int[] nums) {
        // 第一遍读取 stream 将问题复杂度转化为内存可接受的量级.
        int[] bucket = new int[BUCKET_SIZE];
        int step = (right - left) / BUCKET_SIZE;
        boolean even = true;
        int sumCount = 0;
        for (int i = 0; i < nums.length; i++) {
            int index = nums[i] / step;
            bucket[index] = bucket[index] + 1;
            sumCount++;
            even = !even;
        }
        // 如果是偶数,那么就需要计算第 topK 个数
        // 如果是奇数, 那么需要计算第 topK 和 topK+1 的个数.
        int topK = even ? sumCount / 2 : sumCount / 2 + 1;

        int index = 0;
        int indexBucketCount = 0;
        for (index = 0; index < bucket.length; index++) {
            indexBucketCount = bucket[index];
            if (indexBucketCount >= topK) {
                // 当前 bucket 就是中位数的 bucket.
                break;
            }
            topK -= indexBucketCount;
        }
        
        // 划分到这里其实转化为一个 topK 的问题, 再读一遍流.
        if (even && indexBucketCount == topK) { 
            left = index * step;
            right = (index + 2) * step;
            return helperEven(nums, topK);
            // 偶数的时候, 恰好划分到在左右两个子段中.
            // 左右两段中 [topIndex-K + (topIndex-K + 1)] / 2.
        } else if (even) {
            left = index * step;
            right = (index + 1) * step;
            return helperEven(nums, topK);
            // 左边 [topIndex-K + (topIndex-K + 1)] / 2 
        } else {
            left = index * step;
            right = (index + 1) * step;
            return helperOdd(nums, topK);
            // 奇数, 左边 topIndex-K
        }
    }
}

这里边界条件我们处理好之后,关键还是 helperOdd 和 helperEven 这两个函数怎么去求 topK 的问题. 我们还是转化为一个 topK 的问题,那么求 top-K 和 top(K+1)的问题到这里我们是不是可以用堆来解决了呢? 答案是不能,考虑极端情况。 中位数的重复次数非常多

eg:
[100,100,100,100,100...] (1000 亿个 100)

你的划分恰好落到这个桶里面,内存同样会爆掉。

再用时间换空间

假如我们的划分 bucket 大小是 10000,那么最大的时候区间就是 20000。(对应上面的偶数且落到两个分桶的情况) 那么既然划分到某一个 bucket 了,我们直接用数数字的方式来求 topK 就可以了。 我们选用 TreeMap 这种数据结构计数。然后分奇数偶数去求解。

    private double helperEven(int[] nums, int topK) {
        TreeMap<Integer, Integer> map = new TreeMap<>();
        for (int i = 0; i < nums.length; i++) {
            if (nums[i] >= left && nums[i] <= right) {
                if (!map.containsKey(nums[i])) {
                    map.put(nums[i], 1);
                } else {
                    map.put(nums[i], map.get(nums[i]) + 1);
                }
            }
        }

        int count = 0;
        int kNum = Integer.MIN_VALUE;
        int kNextNum = 0;
        for (Map.Entry<Integer, Integer> entry : map.entrySet()) {
            int currentCountIndex = entry.getValue();
            if (kNum != Integer.MIN_VALUE) {
                kNextNum = entry.getKey();
                break;
            }
            if (count + currentCountIndex == topK) {
                kNum = entry.getKey();
            } else if (count + currentCountIndex > topK) {
                kNum = entry.getKey();
                kNextNum = entry.getKey();
                break;
            } else {
                count += currentCountIndex;
            }
        }

        return (kNum + kNextNum) / 2.0;
    }

    private double helperOdd(int[] nums, int topK) {
        TreeMap<Integer, Integer> map = new TreeMap<>();
        for (int i = 0; i < nums.length; i++) {
            if (nums[i] >= left && nums[i] <= right) {
                if (!map.containsKey(nums[i])) {
                    map.put(nums[i], 1);
                } else {
                    map.put(nums[i], map.get(nums[i]) + 1);
                }
            }
        }
        int count = 0;
        int kNum = Integer.MIN_VALUE;
        for (Map.Entry<Integer, Integer> entry : map.entrySet()) {
            int currentCountIndex = entry.getValue();
            if (currentCountIndex + count >= topK) {
                kNum = entry.getKey();
                break;
            } else {
                count += currentCountIndex;
            }
        }

        return kNum;
    }

至此,我觉得算是一个比较好的解决方案，leetcode 社区没有相关的解答和探索，欢迎大家交流。

热门阅读

• 技术文章汇总
• [ Leetcode ] 101. 对称二叉树
• [ Leetcode ] 100. 相同的树
• [ Leetcode ] 98. 验证二叉搜索树

[Acceml]

这是群里一位同学写的，感觉后面可以用大根堆求 topK 没啥问题啊~~

[mnssbe]

群怎么进 拉下我

[pythonee]

是否也拉下我 ZGF5ZGF5dXAxMDI0
另外，好奇的是，这些面试经历都是一个人？一般面试什么岗位会要求当面撕算法？
需要画板当场写吗

[wanderpoet]

求上车+1

[kuangwinnie]

诶 前几天面 tripadvisor 也遇到这题

不过是实习生

[zclHIT]

剑指 offer 上面也有这个题，如果考虑到极致压缩内存的话，真的就有点难了