multiset.gno
5.83 Kb · 213 lines
1// Package multiset is a bag / frequency counter — a set that allows duplicates
2// and remembers how many — as a pure, reusable package.
3//
4// It is the STL multiset and Python's collections.Counter in one type: Add an
5// element several times and the count rises; the distinct elements stay sorted
6// so iteration and rendering are deterministic.
7//
8// The interesting operation is MostCommon(n), and the interesting problem with
9// it is ties. Sorting by count alone leaves elements with equal counts in
10// whatever order the underlying storage happened to yield — which, if that is a
11// built-in map, is unspecified in gno and can differ between nodes. Here the
12// order is total: count descending, then element ascending. Two multisets built
13// from the same elements always produce the same ranking.
14//
15// A live demo of this package is at
16// [r/moul/x/daily/multisetdemo](/r/nym-moultest123/x/daily/multisetdemo/v0).
17package multiset
18
19import "sort"
20
21// MaxDistinct bounds the number of DISTINCT elements so gas stays predictable.
22// Counts themselves are unbounded.
23const MaxDistinct = 4096
24
25// MultiSet counts occurrences of string elements.
26type MultiSet struct {
27 counts map[string]int
28 total int
29}
30
31// New returns an empty MultiSet.
32func New() *MultiSet { return &MultiSet{counts: map[string]int{}} }
33
34// FromSlice builds a MultiSet from elements, counting duplicates.
35func FromSlice(elems []string) *MultiSet {
36 m := New()
37 for _, e := range elems {
38 m.Add(e)
39 }
40 return m
41}
42
43// Add records one occurrence of e. Returns false when e is new and the set
44// already holds MaxDistinct distinct elements.
45func (m *MultiSet) Add(e string) bool { return m.AddN(e, 1) }
46
47// AddN records n occurrences of e. A non-positive n is a no-op returning true.
48func (m *MultiSet) AddN(e string, n int) bool {
49 if n <= 0 {
50 return true
51 }
52 if _, seen := m.counts[e]; !seen && len(m.counts) >= MaxDistinct {
53 return false
54 }
55 m.counts[e] += n
56 m.total += n
57 return true
58}
59
60// Count returns how many times e occurs; zero when absent.
61func (m *MultiSet) Count(e string) int { return m.counts[e] }
62
63// Has reports whether e occurs at least once.
64func (m *MultiSet) Has(e string) bool { return m.counts[e] > 0 }
65
66// Remove drops one occurrence of e, deleting it entirely when the count hits
67// zero. Returns false when e was not present.
68func (m *MultiSet) Remove(e string) bool { return m.RemoveN(e, 1) }
69
70// RemoveN drops up to n occurrences of e. Returns false when e was absent.
71// Removing more than are present clears the element rather than going negative.
72func (m *MultiSet) RemoveN(e string, n int) bool {
73 have, ok := m.counts[e]
74 if !ok || n <= 0 {
75 return false
76 }
77 if n >= have {
78 delete(m.counts, e)
79 m.total -= have
80 return true
81 }
82 m.counts[e] = have - n
83 m.total -= n
84 return true
85}
86
87// RemoveAll drops every occurrence of e. Returns false when e was absent.
88func (m *MultiSet) RemoveAll(e string) bool {
89 have, ok := m.counts[e]
90 if !ok {
91 return false
92 }
93 delete(m.counts, e)
94 m.total -= have
95 return true
96}
97
98// Distinct returns the number of distinct elements.
99func (m *MultiSet) Distinct() int { return len(m.counts) }
100
101// Total returns the sum of every count.
102func (m *MultiSet) Total() int { return m.total }
103
104// IsEmpty reports whether the set holds nothing.
105func (m *MultiSet) IsEmpty() bool { return len(m.counts) == 0 }
106
107// Elements returns the distinct elements, sorted.
108func (m *MultiSet) Elements() []string {
109 out := make([]string, 0, len(m.counts))
110 for e := range m.counts {
111 out = append(out, e)
112 }
113 sort.Strings(out)
114 return out
115}
116
117// Expand returns every occurrence, sorted — a multiset flattened back to a
118// slice. Length equals Total.
119func (m *MultiSet) Expand() []string {
120 out := make([]string, 0, m.total)
121 for _, e := range m.Elements() {
122 for i := 0; i < m.counts[e]; i++ {
123 out = append(out, e)
124 }
125 }
126 return out
127}
128
129// Entry pairs an element with its count.
130type Entry struct {
131 Elem string
132 Count int
133}
134
135// byRank orders entries by count descending, then element ascending — a TOTAL
136// order, so ranking never depends on map iteration.
137type byRank []Entry
138
139func (r byRank) Len() int { return len(r) }
140func (r byRank) Swap(i, j int) { r[i], r[j] = r[j], r[i] }
141func (r byRank) Less(i, j int) bool {
142 if r[i].Count != r[j].Count {
143 return r[i].Count > r[j].Count
144 }
145 return r[i].Elem < r[j].Elem
146}
147
148// MostCommon returns the n most frequent entries, ranked by count descending
149// then element ascending. n <= 0, or larger than the number of distinct
150// elements, returns them all.
151func (m *MultiSet) MostCommon(n int) []Entry {
152 all := make([]Entry, 0, len(m.counts))
153 for e, c := range m.counts {
154 all = append(all, Entry{Elem: e, Count: c})
155 }
156 sort.Sort(byRank(all))
157 if n <= 0 || n > len(all) {
158 return all
159 }
160 return all[:n]
161}
162
163// Union returns a set where each element's count is the MAXIMUM of the two —
164// the standard multiset union.
165func (m *MultiSet) Union(other *MultiSet) *MultiSet {
166 out := m.Clone()
167 for e, c := range other.counts {
168 if c > out.counts[e] {
169 out.setCount(e, c)
170 }
171 }
172 return out
173}
174
175// Intersect returns a set where each element's count is the MINIMUM of the two,
176// keeping only elements present in both.
177func (m *MultiSet) Intersect(other *MultiSet) *MultiSet {
178 out := New()
179 for e, c := range m.counts {
180 if oc, ok := other.counts[e]; ok {
181 if oc < c {
182 c = oc
183 }
184 out.AddN(e, c)
185 }
186 }
187 return out
188}
189
190// Sum returns a set where each element's count is the SUM of the two.
191func (m *MultiSet) Sum(other *MultiSet) *MultiSet {
192 out := m.Clone()
193 for e, c := range other.counts {
194 out.AddN(e, c)
195 }
196 return out
197}
198
199// Clone returns an independent copy.
200func (m *MultiSet) Clone() *MultiSet {
201 out := New()
202 for e, c := range m.counts {
203 out.counts[e] = c
204 }
205 out.total = m.total
206 return out
207}
208
209// setCount overwrites an element's count, keeping total in step.
210func (m *MultiSet) setCount(e string, c int) {
211 m.total += c - m.counts[e]
212 m.counts[e] = c
213}
