ヒープで短い投稿の SNS のタイムラインを設計する

ops = ["PostFeed", "post",
  "post", "getFeed", "follow",
  "post", "getFeed", "unfollow",
  "getFeed"]
args = [[], [1, 101], [2, 201],
  [1], [1, 2], [2, 202], [1],
  [1, 2], [1]]

ops = ["PostFeed", "follow",
  "post", "post", "post",
  "post", "post", "post",
  "post", "post", "post",
  "post", "post", "post",
  "getFeed", "getFeed"]
args = [[], [1, 2],
  [1, 11], [2, 21], [1, 12],
  [2, 22], [1, 13], [2, 23],
  [1, 14], [2, 24], [1, 15],
  [2, 25], [1, 16], [2, 26],
  [1], [2]]

ops = ["PostFeed", "post",
  "follow", "getFeed",
  "unfollow", "getFeed",
  "getFeed"]
args = [[], [5, 50], [5, 5],
  [5], [5, 5], [5], [9]]

from heapq import (
    heapify, heappop, heappush,
)

def subs(sub, u):
    # 初めてなら自分を入れて作る
    if u not in sub:
        sub[u] = {u}
    return sub[u]

class PostFeed:
    def __init__(self):
        self.clock = 0
        self.last = {}
        self.rec = {}
        self.sub = {}

    def post(self, user, pid):
        self.clock += 1
        t = self.clock
        last = self.last
        prev = last.get(user, 0)
        self.rec[t] = pid, prev
        last[user] = t

    def follow(self, a, b):
        if a == b:
            return
        subs(self.sub, a).add(b)

    def unfollow(self, a, b):
        if a == b:
            return
        s = subs(self.sub, a)
        s.discard(b)

    def get_feed(self, user):
        last = self.last
        rec = self.rec
        sub = self.sub
        who = subs(sub, user)
        h = []
        for u in who:
            t = last.get(u, 0)
            if t:
                h.append(-t)
        heapify(h)
        out = []
        for _ in range(10):
            if not h:
                break
            t = -heappop(h)
            pid, prev = rec[t]
            out.append(pid)
            if prev == 0:
                continue
            heappush(h, -prev)
        return out

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

function heapPop(a) {
  const top = a[0];
  const end = a.pop();
  if (!a.length) return top;
  a[0] = end;
  let i = 0;
  for (;;) {
    const l = 2 * i + 1;
    const r = l + 1;
    let m = i;
    if (a[l] > a[m]) m = l;
    if (a[r] > a[m]) m = r;
    if (m === i) return top;
    [a[m], a[i]] = [a[i], a[m]];
    i = m;
  }
}

// 初めてなら自分を入れて作る
function subs(sub, u) {
  if (!sub.has(u)) {
    sub.set(u, new Set([u]));
  }
  return sub.get(u);
}

class PostFeed {
  constructor() {
    this.clock = 0;
    this.last = {};
    this.rec = {};
    this.sub = new Map();
  }

  post(user, pid) {
    this.clock += 1;
    const t = this.clock;
    const last = this.last;
    const prev = last[user];
    this.rec[t] = [pid, prev];
    last[user] = t;
  }

  follow(a, b) {
    if (a === b) return;
    subs(this.sub, a).add(b);
  }

  unfollow(a, b) {
    if (a === b) return;
    const s = subs(this.sub, a);
    s.delete(b);
  }

  getFeed(user) {
    const { rec, sub } = this;
    const who = subs(sub, user);
    const h = [];
    for (const u of who) {
      const t = this.last[u];
      if (t) heapPush(h, t);
    }
    const out = [];
    while (out.length < 10) {
      if (h.length === 0) break;
      const t = heapPop(h);
      const [id, prev] = rec[t];
      out.push(id);
      if (prev > 0) {
        heapPush(h, prev);
      }
    }
    return out;
  }
}

import "container/heap"

type maxHeap []int

func (h maxHeap) Len() int {
	return len(h)
}

func (h maxHeap) Less(
	i, j int,
) bool {
	return h[i] > h[j]
}

func (h maxHeap) Swap(
	i, j int,
) {
	h[i], h[j] = h[j], h[i]
}

func (h *maxHeap) Push(x any) {
	*h = append(*h, x.(int))
}

func (h *maxHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[:n-1]
	return x
}

type record struct {
	id, prev int
}

type userSet map[int]bool

type PostFeed struct {
	clock int
	last  map[int]int
	rec   map[int]record
	sub   map[int]userSet
}

func newPostFeed() *PostFeed {
	return &PostFeed{
		last: map[int]int{},
		rec:  map[int]record{},
		sub:  map[int]userSet{},
	}
}

// 初めてなら自分を入れて作る
func subs(
	sub map[int]userSet, u int,
) userSet {
	s, ok := sub[u]
	if !ok {
		s = userSet{u: true}
		sub[u] = s
	}
	return s
}

func (f *PostFeed) Post(
	user, id int,
) {
	f.clock++
	t := f.clock
	prev := f.last[user]
	f.rec[t] = record{id, prev}
	f.last[user] = t
}

func (f *PostFeed) Follow(
	a, b int,
) {
	if a != b {
		subs(f.sub, a)[b] = true
	}
}

func (f *PostFeed) Unfollow(
	a, b int,
) {
	if a != b {
		delete(subs(f.sub, a), b)
	}
}

func (f *PostFeed) GetFeed(
	user int,
) []int {
	h := &maxHeap{}
	who := subs(f.sub, user)
	for u := range who {
		if t := f.last[u]; t > 0 {
			*h = append(*h, t)
		}
	}
	heap.Init(h)
	out := []int{}
	for h.Len() > 0 {
		if len(out) == 10 {
			break
		}
		t := heap.Pop(h).(int)
		r := f.rec[t]
		out = append(out, r.id)
		if r.prev > 0 {
			heap.Push(h, r.prev)
		}
	}
	return out
}