Top-k

12. 11. 2020

Ve většině programovacích jazyků mi chybí top-k operace přímo ve standardní knihovně. Představte si situaci, kdy máte určitou kolekci a chcete z ní dostat 10 elementů s největší váhou dle určité funkce. Dosáhnout se toho dá celkem pohodlně řazením, ale to má jeden zásadní nedostatek: Je neefektivní a pomalé.

case class Item(id: Int, weight: Double)
val items: Seq[Item] = ???

xs.sortBy(-_.weight).take(k)

Mnohokrát se mi stalo, že jsem si říkal, seřadím to, přece to nemůže být tak pomalé, jen aby se později ukázalo, že v řazení program tráví většinu času. Řazení běží v čase O(n log(n)), ale už skoro šedesát let známe quickselect nebo heap, které to samé zvládnou v čase O(n) respektive O(n log(k)).

Ideálně by mělo jít být oneliner, který by byl nejen stručný, ale i efektivní.

xs.topk(10)(_.weight)

V Javě, Scale a dalších jazycích odkojených JVM se dá použít java.util.SortedSet nebo scala.collection.mutable.SortedSet.

val topk = scala.collection.mutable.TreeSet[Item]()(Ordering.by(_.weight))
val k = 10

for (item <- items) {
  topk += item
  if (topk.size > k) topk.remove(topk.head)
}

println(topk)

To je jednak nepříjemně ukecané, ale hlavě má jeden zásadní problém: TreeSet nemůže obsahovat duplicity. Ne duplicity podle equals, které dávají smysl, ale duplicity podle řazení. V dokumentaci se to označuje jako konzistentní s equals.

The natural ordering for a class C is said to be consistent with equals if and only if e1.compareTo(e2) == 0 has the same boolean value as e1.equals(e2) for every e1 and e2 of class C.

Item(id = 47, weight = 1.0) je považován za stejný jako Item(id = 69696969, weight = 1.0) při řazení podle weight a výsledná množina bude obsahovat jen jeden z nich, což je přesně opak intuitivního chování. Člověk by je chtěl oba a vyhazovat jen duplicity, které a.equals(b). TreeSet ale bere jako duplicity všechno a.weight.equals(b.weight).

Správná třída je PriorityQueue. Ta konečně dělá, co chci. Stále ne moc elegantně, ale aspoň správně.

val k = 10
val heap = new java.util.PriorityQueue[Item](1, Ordering.by(_.weight))

for (item <- items) {
  heap.add(item)
  if (heap.size > k) heap.poll()
}

val topk = Seq.fill(heap.size) { heap.poll() }.reverse

Navíc tady je ta komplikace, že API nemá metodu, jak en-masse vytahat elementy seřazené a musím to udělat ručně, což to má stále daleko od pohodlného one-lineru. Javadoc radí Arrays.sort(heap.toArray()), ale řazení je přesně to, čemu se chci vyhnout.

Naproti tomu jazyk D, se kterým se v posledních dnech intimně seznamuji, má ve standardní knihovně funkci topN pro materializované kolekce, topNCopy pro iterátory a BinaryHeap s příjemným rozhraním pro všechny ostatní případy.

struct Sim { uint id; float weight; }

auto heap = BinaryHeap!(Sim[], "a.weight > b.weight")(new Sim[30], 0);
foreach (item; items) {
  heap.conditionalInsert(Sim(item.id, computeWeight(item));
}

Metoda conditionalInsert dělá přesně co, potřebuji: Pokud je v haldě místo, element vloží, pokud ne, vloží ho tam jen, když je větší než současné minimum.

auto result = new Sim[30];

items.map!((item) => Sim(item.id, computeWeight(item)))
  .topNCopy!"a.weight > b.weight"(result);

Takhle bych si to představoval.

D je super v mnoha dalších ohledech. Autoři si dávají velice záležet na efektivitě² výsledného kódu i efektivitě použité implementace. Jako jeden z mála jazyků přímo ve standardní knihovna nabízí schwartzSort, což je řazení za pomoci Schwartzian transform pro případy, kdy je porovnávací funkce příliš drahá. Schwartz sort člověk nepoužije příliš často, ale když ho potřebuje, potřebuje ho jako sůl.

0xDEADBEEF

Top-k