1
package mergeiter
2

3
import "container/heap"
4

5
// 堆元素：当前值 + 来源流编号
6
type entry struct {
7
    value int
8
    index int // 属于哪个流
9
}
10

11
// 最小堆
12
type minHeap []entry
13

14
func (h minHeap) Len() int            { return len(h) }
15
func (h minHeap) Less(i, j int) bool  { return h[i].value < h[j].value }
16
func (h minHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
17
func (h *minHeap) Push(x any)         { *h = append(*h, x.(entry)) }
18
func (h *minHeap) Pop() any {
19
    old := *h
20
    n := len(old)
21
    v := old[n-1]
22
    *h = old[:n-1]
23
    return v
24
}
25

26
// 合并迭代器
27
type MergeIterator struct {
28
    heap    minHeap
29
    streams [][]int // K 路有序流
30
    cursors []int   // 每个流的当前位置
31
}
32

33
func New(streams [][]int) *MergeIterator {
34
    it := &MergeIterator{
35
        streams: streams,
36
        cursors: make([]int, len(streams)),
37
        heap:    make(minHeap, 0, len(streams)),
38
    }
39
    // 各流首元素入堆
40
    for i, s := range streams {
41
        if len(s) > 0 {
42
            heap.Push(&it.heap, entry{value: s[0], index: i})
43
            it.cursors[i] = 1
44
        }
45
    }
46
    heap.Init(&it.heap)
47
    return it
48
}
49

50
// 推进到下一个元素，返回是否还有数据
51
func (it *MergeIterator) Next() (int, bool) {
52
    if it.heap.Len() == 0 {
53
        return 0, false
54
    }
55
    // 弹出全局最小值
56
    top := heap.Pop(&it.heap).(entry)
57
    // 从该流补充下一个元素
58
    if it.cursors[top.index] < len(it.streams[top.index]) {
59
        heap.Push(&it.heap, entry{
60
            value: it.streams[top.index][it.cursors[top.index]],
61
            index: top.index,
62
        })
63
        it.cursors[top.index]++
64
    }
65
    return top.value, true
66
}

1
// 最小堆元素
2
interface Entry {
3
  value: number;
4
  streamIdx: number;
5
}
6

7
// 最小堆
8
class MinHeap {
9
  private data: Entry[] = [];
10

11
  get size(): number {
12
    return this.data.length;
13
  }
14

15
  push(e: Entry): void {
16
    this.data.push(e);
17
    this.bubbleUp(this.data.length - 1);
18
  }
19

20
  pop(): Entry | undefined {
21
    if (this.data.length === 0) return undefined;
22
    const top = this.data[0];
23
    const last = this.data.pop()!;
24
    if (this.data.length > 0) {
25
      this.data[0] = last;
26
      this.sinkDown(0);
27
    }
28
    return top;
29
  }
30

31
  private bubbleUp(i: number): void {
32
    while (i > 0) {
33
      const parent = (i - 1) >> 1;
34
      if (this.data[i].value >= this.data[parent].value) break;
35
      [this.data[i], this.data[parent]] = [this.data[parent], this.data[i]];
36
      i = parent;
37
    }
38
  }
39

40
  private sinkDown(i: number): void {
41
    const n = this.data.length;
42
    while (true) {
43
      let smallest = i;
44
      const left = 2 * i + 1, right = 2 * i + 2;
45
      if (left < n && this.data[left].value < this.data[smallest].value) smallest = left;
46
      if (right < n && this.data[right].value < this.data[smallest].value) smallest = right;
47
      if (smallest === i) break;
48
      [this.data[i], this.data[smallest]] = [this.data[smallest], this.data[i]];
49
      i = smallest;
50
    }
51
  }
52
}
53

54
// 合并迭代器
55
class MergeIterator {
56
  private heap: MinHeap = new MinHeap();
57
  private cursors: number[];
58

59
  constructor(private streams: number[][]) {
60
    this.cursors = new Array(streams.length).fill(0);
61
    // 各流首元素入堆
62
    for (let i = 0; i < streams.length; i++) {
63
      if (streams[i].length > 0) {
64
        this.heap.push({ value: streams[i][0], streamIdx: i });
65
        this.cursors[i] = 1;
66
      }
67
    }
68
  }
69

70
  next(): { value: number; done: false } | { done: true } {
71
    if (this.heap.size === 0) return { done: true };
72
    const top = this.heap.pop()!;
73
    // 从被弹出的流补充下一个元素
74
    const stream = this.streams[top.streamIdx];
75
    if (this.cursors[top.streamIdx] < stream.length) {
76
      this.heap.push({
77
        value: stream[this.cursors[top.streamIdx]],
78
        streamIdx: top.streamIdx,
79
      });
80
      this.cursors[top.streamIdx]++;
81
    }
82
    return { value: top.value, done: false };
83
  }
84
}