2 本のヒープで流れてくる数の中央値をすぐ答える仕組みを作る｜コーディング面接問題

問題

計測の値が 1 つずつ届く場面で、届いた値の中央値をいつでも答えられる仕組みを、クラス StreamMedian として作ってください。次の 2 つの操作を持ちます。

add(num) は、整数 num を足します。何も返しません
median() は、それまでに足した数の中央値を返します

中央値は、足した数を小さい順に並べたときの真ん中の値です。数の個数が偶数のときは、真ん中の 2 つを足して 2 で割った値です。小数になることがあり、そのときは 2.5 のような数で返します。割り切れるときは 3 と返しても 3.0 と返してもかまいません。

メソッドの名前は、JavaScript と Python は上のまま、Go は Add、Median と頭を大文字にします。Go では newStreamMedian() で作り、Median は float64 を返します。

入力は、呼ぶ操作の名前の並び ops と、それぞれの引数の並び args で渡します。ops の最初はクラスの名前で、args の最初は空の並びです。出力は操作ごとの戻り値の並びで、作った所と add は null です。

median は、数が 1 つ以上足されたあとにだけ呼びます
num は -100000 以上 100000 以下です
操作の回数は 50000 回以下で、add と median がどんな順で混ざってもかまいません

入力と出力の例

例 1

4 だけなら 4、4 と 10 なら真ん中の 2 つを足して 2 で割って 7、1 を足すと並びは 1、4、10 で 4 です。

入力

ops = ["StreamMedian", "add",
  "median", "add", "median",
  "add", "median"]
args = [[], [4], [], [10], [],
  [1], []]

出力

[null, null, 4, null, 7,
 null, 4]

例 2

3 と 6 の中央値は 4.5 です。同じ値も別々に数えます。

入力

ops = ["StreamMedian", "add",
  "add", "median", "add",
  "median", "add", "median"]
args = [[], [3], [6], [], [6],
  [], [-2], []]

出力

[null, null, null, 4.5, null,
 6, null, 4.5]

例 3

負の数が小さくなる向きに届いても、同じように答えます。

入力

ops = ["StreamMedian", "add",
  "add", "median", "add",
  "median", "add", "median"]
args = [[], [-1], [-5], [],
  [-9], [], [-20], []]

出力

[null, null, null, -3, null,
 -5, null, -7]

面接での聞かれ方

面接では、応答時間の監視画面のように、値が流れ続ける場面を例に出して、仕組みの設計として渡されることが多いです。

「add と median が同じくらいの回数で呼ばれます。どちらも速くしてください」と、両方の重さを条件にされます
書き終えたあとに「値がすべて 0 以上 100 以下だと分かっていたら、どう持ちますか？」「ほとんどがその範囲で、ときどき外れる値が来るなら？」と、値の分布に合わせた持ち方を聞かれます
「直近の 1000 個だけの中央値にするなら、古い値をどう消しますか？」「機械を何台かに分けるなら？」と、仕組みを広げた先まで話が進みます

テックリードとしては、1 つの書き方を書き切るだけでなく、「どの操作が何回呼ばれるか」「値がどんな範囲に来るか」で持ち方を選び直せることが見られています。

アルゴリズムのチートシートの「ハッシュマップ」に、この型の要点を表とコードでまとめています。

まず考える手順

段 1

まずは素直に、足した数を並びに入れておき、median のたびに並べ替えて真ん中を見る書き方を考えます。add は O(1) ですが、median は毎回 O(n log n) で、何度も聞かれると重くなります。

段 2

並びをいつも小さい順に保つ書き方も考えます。二分探索で入れる位置は O(log n) で見つかりますが、途中に入れるには後ろをずらすので add が O(n) です。median は真ん中を読むだけの O(1) になりますが、add の重さが残ります。

段 3

中央値に要るのは、小さい順に並べたときの真ん中の 1 つか 2 つだけで、ほかの数の並び順は要りません。そこで、数を「小さい半分」と「大きい半分」に分け、小さい半分はいちばん大きい値をすぐ出せるヒープ (最大ヒープ)、大きい半分はいちばん小さい値をすぐ出せるヒープ (最小ヒープ) に入れます。2 つのヒープの先頭が、ちょうど真ん中の境目です。

段 4

2 つの約束を守ります。1 つ目は、小さい半分のどの数も、大きい半分のどの数以下であること。2 つ目は、小さい半分の個数が大きい半分と同じか、1 つだけ多いことです。足すときは、まず小さい半分に入れ、小さい半分のいちばん大きい数を大きい半分へ移します。これで 1 つ目の約束が守られます。そのあと大きい半分のほうが多くなっていたら、大きい半分のいちばん小さい数を小さい半分へ戻します。これで 2 つ目も守られます。

段 5

median は、個数が奇数なら小さい半分の先頭、偶数なら 2 つの先頭を足して 2 で割った値です。add はヒープの出し入れが 3 回で O(log n)、median は先頭を見るだけで O(1) です。

段 6

ヒープを持たない言語では、並びの上にヒープを書きます。親の位置は「自分の番号から 1 を引いて半分にした位置」、子は「自分の番号の 2 倍に 1 か 2 を足した位置」です。最大ヒープは、値の符号を逆にして最小ヒープに入れれば、同じ関数で済みます。

解答

Python

from heapq import heappop
from heapq import heappush


class StreamMedian:
    def __init__(self):
        self.lo = []
        self.hi = []

    def add(self, num):
        lo = self.lo
        hi = self.hi
        heappush(lo, -num)
        x = -heappop(lo)
        heappush(hi, x)
        if len(hi) > len(lo):
            y = heappop(hi)
            heappush(lo, -y)

    def median(self):
        lo = self.lo
        hi = self.hi
        a = -lo[0]
        if len(lo) > len(hi):
            return float(a)
        b = hi[0]
        return (a + b) / 2

JavaScript

class StreamMedian {
  constructor() {
    this.lo = [];
    this.hi = [];
  }

  add(num) {
    const lo = this.lo;
    const hi = this.hi;
    heapPush(lo, -num);
    const x = -heapPop(lo);
    heapPush(hi, x);
    if (hi.length > lo.length) {
      const y = heapPop(hi);
      heapPush(lo, -y);
    }
  }

  median() {
    const lo = this.lo;
    const hi = this.hi;
    const a = -lo[0];
    if (lo.length > hi.length) {
      return a;
    }
    return (a + hi[0]) / 2;
  }
}

function heapPush(h, x) {
  h.push(x);
  let i = h.length - 1;
  while (i > 0) {
    const p = (i - 1) >> 1;
    if (h[p] <= h[i]) {
      break;
    }
    [h[p], h[i]] = [h[i], h[p]];
    i = p;
  }
}

function heapPop(h) {
  const top = h[0];
  const last = h.pop();
  const n = h.length;
  if (n === 0) {
    return top;
  }
  h[0] = last;
  let i = 0;
  while (true) {
    let m = i;
    const l = 2 * i + 1;
    const r = l + 1;
    if (l < n && h[l] < h[m]) {
      m = l;
    }
    if (r < n && h[r] < h[m]) {
      m = r;
    }
    if (m === i) {
      return top;
    }
    [h[m], h[i]] = [h[i], h[m]];
    i = m;
  }
}

Go

type StreamMedian struct {
	lo  []int
	hi []int
}

func newStreamMedian(
) *StreamMedian {
	return &StreamMedian{}
}

func (s *StreamMedian) Add(
	num int,
) {
	heapPush(&s.lo, -num)
	x := -heapPop(&s.lo)
	heapPush(&s.hi, x)
	if len(s.hi) > len(s.lo) {
		y := heapPop(&s.hi)
		heapPush(&s.lo, -y)
	}
}

func (s *StreamMedian) Median(
) float64 {
	a := -s.lo[0]
	if len(s.lo) > len(s.hi) {
		return float64(a)
	}
	b := s.hi[0]
	return float64(a+b) / 2
}

func heapPush(h *[]int, x int) {
	a := append(*h, x)
	i := len(a) - 1
	for i > 0 {
		p := (i - 1) / 2
		if a[p] <= a[i] {
			break
		}
		a[p], a[i] = a[i], a[p]
		i = p
	}
	*h = a
}

func heapPop(h *[]int) int {
	a := *h
	top := a[0]
	n := len(a) - 1
	a[0] = a[n]
	a = a[:n]
	i := 0
	for {
		m := i
		l := 2*i + 1
		r := l + 1
		if l < n && a[l] < a[m] {
			m = l
		}
		if r < n && a[r] < a[m] {
			m = r
		}
		if m == i {
			break
		}
		a[m], a[i] = a[i], a[m]
		i = m
	}
	*h = a
	return top
}

lo には符号を逆にした数を入れているので、lo の先頭の符号を戻した値が、小さい半分のいちばん大きい数です。Python は備わった heapq を使い、JavaScript と Go は、クラスの外に置いた heapPush と heapPop で並びをヒープとして扱います。

計算量

時間は、add が O(log n)、median が O(1) です。add はヒープへの出し入れを最大 3 回し、それぞれ木の高さの O(log n) です
使う記憶は、足した数の個数を n として O(n) です。どの数も 2 本のヒープのどちらか一方に 1 回だけ入ります
値の範囲が 0 以上 100 以下のように小さく決まっているなら、値ごとの個数を数える 101 個の表で持つ手もあります。add は O(1)、median は表を前から数えて O(101) で、値の範囲が狭いほど得です。範囲の外の値がまれに来るなら、外れた値だけ別のヒープや並びに分けて持ち、表と組み合わせます

よくある間違い

新しい数を、個数の少ないほうのヒープへそのまま入れる。個数はそろいますが、2 つの半分の境目が崩れます。例 3 で -1 のあとの -5 を大きい半分へ入れると、小さい半分の -1 より小さい数が大きい半分に入り、-9 を足したあとの中央値が -5 ではなく -1 になります。いったん反対のヒープを通してから入れます
偶数のときの割り算を整数の割り算にする。 Python の // や、Go の整数どうしの割り算では、4.5 が 4 になります。負の数ではさらに外れ、-1 と -2 の中央値 -1.5 が、Python の // では -2、Go の整数の割り算では -1 になります。小数で割ります
最大ヒープの符号の戻し忘れ。小さい半分には符号を逆にして入れているので、先頭をそのまま使うと、中央値の符号が逆になります。取り出すときと見るときは、必ず符号を戻します

解答を動かした版 Python 3.12.14Node.js 22.23.4Go 1.24.6確認日 2026年10月30日

class StreamMedian: def __init__(self): # ここに持つ物を書きます pass def add(self, num): # ここに書きます pass def median(self): # ここに書きます pass # 「実行」を押すと、下の行が動きます obj = StreamMedian() print(obj.add(4))

2 本のヒープで流れてくる数の中央値をすぐ答える仕組みを作る

問題