// Copyright 2012 Stefan Nilsson // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // Package bit provides a bitset implementation and utility bit functions for uint64 words. package bit import ( "bytes" "strconv" ) const ( bpw = 64 // bits per word maxw = 1<>shift] & (1<<(n&mask)) == 1 iff n belongs to set, // • data[i] == 0 for all i such that len(data) ≤ i < cap(data), // • len(data) == 0 if set is empty, // • data[len(data)-1] != 0 if set is nonempty, // • min is the minimum element of a nonempty set, 0 otherwise. data []uint64 min int } // New creates a new set S with the given elements. func New(n ...int) (S *Set) { if len(n) == 0 { return new(Set) } min, max := n[0], n[0] for _, e := range n { switch { case e > max: max = e case e < min: min = e } } S = &Set{ data: make([]uint64, max>>shift+1), min: min, } for _, e := range n { S.data[e>>shift] |= 1 << uint(e&mask) } return } // Add adds n to S, setting S to S ∪ {n}, and returns the updated set. func (S *Set) Add(n int) *Set { if n < 0 { panic("Set: Add(" + strconv.Itoa(n) + ") index out of range") } if len(S.data) == 0 || n < S.min { S.min = n } i := n >> shift if i >= len(S.data) { S.resize(i + 1) } S.data[i] |= 1 << uint(n&mask) return S } // AddRange adds all integers between m and n-1, m ≤ n, // setting S to S ∪ {m..n-1}, and returns the updated set. func (S *Set) AddRange(m, n int) *Set { if m == n && m >= 0 { return S } if m > n || m < 0 { panic("Set: AddRange(" + strconv.Itoa(m) + ", " + strconv.Itoa(n) + ") bounds out of range") } if len(S.data) == 0 || m < S.min { S.min = m } n-- low, high := m>>shift, n>>shift if high >= len(S.data) { S.resize(high + 1) } d := S.data if low == high { // Range fits in one word. d[low] |= bitMask(m&mask, n&mask) } else { // Range spans at least two words. d[low] |= bitMask(m&mask, bpw-1) for i := low + 1; i < high; i++ { d[i] = maxw } d[high] |= bitMask(0, n&mask) } return S } // Remove removes n from S, setting S to S ∖ {n}, and returns the updated set. func (S *Set) Remove(n int) *Set { if n < 0 { panic("Set: Remove(" + strconv.Itoa(n) + ") index out of range") } if n < S.min { return S } i := n >> shift if i >= len(S.data) { return S } S.data[i] &^= 1 << uint(n&mask) S.trim() if n == S.min { S.min = findMinFrom(S.min, S.data) } return S } // RemoveMin removes the minimum element from S, setting S to S ∖ {min}, // and returns min. It panics if S is empty. func (S *Set) RemoveMin() (min int) { if len(S.data) == 0 { panic("Set: RemoveMin not defined for empty set") } min = S.min i := min >> shift d := S.data d[i] &^= 1 << uint(min&mask) if d[i] != 0 || i+1 < len(d) { // There are more elements left. S.min = findMinFrom(min, d) } else { S.Clear() // The array contains only zeroes. } return } // RemoveMax removes the maximum element from S, setting S to S ∖ {max}, // and returns max. It panics if S is empty. func (S *Set) RemoveMax() (max int) { if len(S.data) == 0 { panic("Set: RemoveMax not defined for empty set") } d := S.data i := len(d) - 1 w := d[i] max = i<= max { return } i = m >> shift s := 1 + uint(m&mask) w := d[i] >> s << s // Zero out bits for numbers ≤ m. for w == 0 { i++ w = d[i] } return i< max { return max, true } i = m >> shift s := bpw - uint(m&mask) w := d[i] << s >> s // Zero out bits for numbers ≥ m. for w == 0 { i-- w = d[i] } return i<= 0 { return S } if m > n || m < 0 { panic("Set: RemoveRange(" + strconv.Itoa(m) + ", " + strconv.Itoa(n) + ") bounds out of range") } n-- d := S.data low, high := m>>shift, n>>shift // Range does not intersect S. if n < S.min || low >= len(d) { return S } // Bottom of range undershoots S. if m < S.min { low = S.min >> shift // low ≤ high still holds, since S.min ≤ n. m = 0 // To assure that m&mask == 0 below. } // Top of range overshoots S. if len(d) <= high { high = len(d) - 1 // low ≤ high still holds, since low < len(d). n = bpw - 1 // To assure that n&mask == bpw-1 below. } if low == high { // Range fits in one word d[low] &^= bitMask(m&mask, n&mask) } else { // Range spans at least two words d[low] &^= bitMask(m&mask, bpw-1) for i := low + 1; i < high; i++ { d[i] = 0 } d[high] &^= bitMask(0, n&mask) } S.trim() if m <= S.min { S.min = findMinFrom(S.min, S.data) } return S } // Clear removes all elements and returns the updated empty set. func (S *Set) Clear() *Set { S.realloc(0) S.min = 0 return S } // Flip removes n from S if it is present, otherwise it adds n, // setting S to S ∆ {n}, and returns the updated set. func (S *Set) Flip(n int) *Set { if n < 0 { panic("Set: Flip(" + strconv.Itoa(n) + ") index out of range") } if len(S.data) == 0 { // Set correct minimum when S empty. S.min = n } i := n >> shift if i >= len(S.data) { S.resize(i + 1) } S.data[i] ^= 1 << uint(n&mask) S.trim() if n <= S.min { // If flip removed S.min, we must search for a new minimum. // If n < S.min, the search will find n right away. S.min = findMinFrom(n, S.data) } return S } // FlipRange flips all elements in the range m..n-1, m ≤ n, // setting S to S ∆ {m..n-1}, and returns the updated set. func (S *Set) FlipRange(m, n int) *Set { if m == n && m >= 0 { return S } if m > n || m < 0 { panic("Set: FlipRange(" + strconv.Itoa(m) + ", " + strconv.Itoa(n) + ") bounds out of range") } d := S.data if len(d) == 0 { // Set correct minimum when S empty. S.min = m } n-- low, high := m>>shift, n>>shift if high >= len(d) { S.resize(high + 1) d = S.data } if low == high { // Range fits in one word. d[low] ^= bitMask(m&mask, n&mask) } else { // Range spans at least two words. d[low] ^= bitMask(m&mask, bpw-1) for i := low + 1; i < high; i++ { d[i] ^= maxw } d[high] ^= bitMask(0, n&mask) } S.trim() // If S.min < m, the minimum does not change. if m <= S.min { // The new minimum can't be smaller than m. S.min = findMinFrom(m, S.data) } return S } // Contains returns true if n, n ≥ 0, is an element of S. func (S *Set) Contains(n int) bool { d := S.data i := n >> shift if i >= len(d) { return false } return d[i]&(1<>shift, len(d); i < l; i++ { if w := d[i]; w != 0 { n += Count(w) } } return n } // IsEmpty returns true if S = ∅. func (S *Set) IsEmpty() bool { return len(S.data) == 0 } // Equals returns true if A and B contain the same elements. func (A *Set) Equals(B *Set) bool { if A == B { return true } if A.min != B.min { return false } a, b := A.data, B.data if len(a) != len(b) { return false } for i, l := A.min>>shift, len(a); i < l; i++ { if a[i] != b[i] { return false } } return true } // SubsetOf returns true if A ⊆ B. func (A *Set) SubsetOf(B *Set) bool { if A == B { return true } a, b := A.data, B.data if len(a) > len(b) { return false } if len(a) != 0 && A.min < B.min { return false } for i, l := A.min>>shift, len(a); i < l; i++ { if a[i]&^b[i] != 0 { return false } } return true } // Intersects returns true if A and B overlap, i.e. A ∩ B ≠ ∅. func (A *Set) Intersects(B *Set) bool { a, b := A.data, B.data if len(a) == 0 || len(b) == 0 { return false } if A == B { // Notice that both sets are nonempty. return true } m := min(len(a), len(b)) for i := max(A.min, B.min) >> shift; i < m; i++ { if a[i]&b[i] != 0 { return true } } return false } // And creates a new intersection S = A ∩ B that consists of all // elements that belong to both A and B. func (A *Set) And(B *Set) (S *Set) { return new(Set).SetAnd(A, B) } // Or creates a new union S = A ∪ B that contains all // elements that belong to either A or B. func (A *Set) Or(B *Set) (S *Set) { return new(Set).SetOr(A, B) } // AndNot creates a new set difference S = A ∖ B that // consists of all elements that belong to A, but not to B. func (A *Set) AndNot(B *Set) (S *Set) { return new(Set).SetAndNot(A, B) } // Xor creates a new symmetric difference S = A ∆ B = (A ∪ B) ∖ (A ∩ B) // that consists of all elements that belong to either A or B, but not to both. func (A *Set) Xor(B *Set) (S *Set) { return new(Set).SetXor(A, B) } // SetWord interprets w as a bitset with numbers in the range 64i to 64i + 63, // where 0 ≤ i ≤ ⌊MaxInt/64⌋, overwrites this range in S with w, and returns S. func (S *Set) SetWord(i int, w uint64) *Set { if i > MaxInt/64 { panic("Set: SetWord(" + strconv.Itoa(i) + ", 0x" + strconv.FormatUint(w, 16) + ") index out of range") } if w == 0 && i >= len(S.data) { return S } if len := len(S.data); i >= len { if len == 0 { // Set tentative minimum when S is empty. S.min = 64 * i } S.resize(i + 1) } S.data[i] = w if w == 0 && i == len(S.data)-1 { S.trim() } if mi := S.min >> shift; i < mi && w != 0 || i == mi { S.min = findMinFrom(64*i, S.data) } return S } // Word returns the range 64i to 64i + 63 of S as a bitset represented by w. func (S *Set) Word(i int) (w uint64) { if i < len(S.data) { w = S.data[i] } return } // Set sets S to A and returns S. func (S *Set) Set(A *Set) *Set { S.realloc(len(A.data)) low := min(A.min>>shift, S.min>>shift) // Words below this index are all zero, copy(S.data[low:], A.data[low:]) // so there is no need to copy them. S.min = A.min return S } // SetAnd sets S to the intersection A ∩ B and returns S. func (S *Set) SetAnd(A, B *Set) *Set { a, b := A.data, B.data la, lb := len(a), len(b) ma, mb, ms := A.min, B.min, S.min // Check if ranges overlap. if la == 0 || lb == 0 || la <= mb>>shift || lb <= ma>>shift { return S.Clear() } // S.data[i], i < low, are words that should be zero and already are. low := min(max(ma, mb), ms) >> shift // Find last nonzero word in result. i := min(la, lb) - 1 for i >= 0 && a[i]&b[i] == 0 { i-- } if S == A || S == B { S.resize(i + 1) } else { S.realloc(i + 1) } for ; i >= low; i-- { S.data[i] = a[i] & b[i] } S.min = findMinFrom(max(ma, mb), S.data) return S } // SetAndNot sets S to the set difference A ∖ B and returns S. func (S *Set) SetAndNot(A, B *Set) *Set { a, b := A.data, B.data la, lb := len(a), len(b) ma, mb, ms := A.min, B.min, S.min // Check if ranges overlap. if la == 0 || lb == 0 || la <= mb>>shift || lb <= ma>>shift { return S.Set(A) } // S.data[i], i < low, are words that should be zero and already are. low := min(ma, ms) >> shift // Result requires len(a) words if len(a) > len(b), // otherwise find last nonzero word in result. i := la - 1 if la <= lb { for i >= 0 && a[i]&^b[i] == 0 { i-- } } if S == A || S == B { S.resize(i + 1) } else { S.realloc(i + 1) } d := S.data if m := max(lb, low); m <= i { copy(d[m:i+1], a[m:i+1]) i = m - 1 } for ; i >= low; i-- { d[i] = a[i] &^ b[i] } if A.min < B.min { // Minimum can't have changed. S.min = A.min } else { S.min = findMinFrom(ma, d) // can't be smaller than previous A.min. } return S } // SetOr sets S to the union A ∪ B and returns S. func (S *Set) SetOr(A, B *Set) *Set { // Swap, if necessary, to make A shorter than B. if len(A.data) > len(B.data) { A, B = B, A } a, b := A.data, B.data la, lb := len(a), len(b) ma, mb, ms := A.min, B.min, S.min i := lb - 1 if S == A || S == B { S.resize(i + 1) } else { S.realloc(i + 1) } d := S.data copy(d[la:i+1], b[la:i+1]) i = la - 1 // S.data[i], i < low, are words that should be zero and already are. low := min(min(ma, mb), ms) >> shift for ; i >= low; i-- { d[i] = a[i] | b[i] } if la == 0 { S.min = mb } else if lb == 0 { S.min = ma } else { S.min = min(ma, mb) } return S } // SetXor sets S to the symmetric difference A ∆ B = (A ∪ B) ∖ (A ∩ B) // and returns S. func (S *Set) SetXor(A, B *Set) *Set { // Swap, if necessary, to make A shorter than B. if len(A.data) > len(B.data) { A, B = B, A } a, b := A.data, B.data la, lb := len(a), len(b) ma, mb, ms := A.min, B.min, S.min i := lb - 1 if la == lb { // The only case where result may be shorter than len(b). low := min(ma, mb) >> shift for i >= low && a[i]^b[i] == 0 { i-- } if i == low-1 { // Only zero elements left. return S.Clear() } } if S == A || S == B { S.resize(i + 1) } else { S.realloc(i + 1) } d := S.data if la <= i { copy(d[la:i+1], b[la:i+1]) i = la - 1 } // S.data[i], i < low, are words that should be zero and already are. low := min(min(ma, mb), ms) >> shift for ; i >= low; i-- { d[i] = a[i] ^ b[i] } S.min = findMinFrom(min(ma, mb), d) return S } // Do calls function f for each element n ∊ S in numerical order. // It is safe for f to add or remove elements e, e ≤ n, from S. // The behavior of Do is undefined if f changes the set in any other way. func (S *Set) Do(f func(n int)) { d := S.data for i, l := S.min>>shift, len(d); i < l; i++ { w := d[i] if w == 0 { continue } n := i << shift // element represented by w&1 for w != 0 { b := MinPos(w) n += b f(n) n++ w >>= uint(b + 1) for w&1 != 0 { // common case f(n) n++ w >>= 1 } } } } // String returns a string representation of S. // The elements are listed in ascending order, enclosed by braces, // and separated by ", ". Runs of at least three elements a, a+1, ..., b // are given as "a..b". // // For example, the set {1, 2, 6, 5, 3} is represented as "{1..3, 5, 6}". func (S *Set) String() string { sb := new(bytes.Buffer) sb.WriteString("{") a, b := -1, -2 // keeps track of a range a..b of elements S.Do(func(n int) { if n == b+1 { // Increase current range from a..b to a..b+1. b = n return } writeRange(sb, a, b) a, b = n, n // Start new range. }) writeRange(sb, a, b) if S.Size() > 0 { sb.Truncate(sb.Len() - 2) // Remove trailing ", ". } sb.WriteString("}") return sb.String() } // writeRange writes either "", "a", "a, b, " or "a..b, " to buffer. func writeRange(sb *bytes.Buffer, a, b int) { switch { case a > b: // sb.WriteString(sb, "") case a == b: sb.WriteString(strconv.Itoa(a)) sb.WriteString(", ") case a+1 == b: sb.WriteString(strconv.Itoa(a)) sb.WriteString(", ") sb.WriteString(strconv.Itoa(b)) sb.WriteString(", ") default: sb.WriteString(strconv.Itoa(a)) sb.WriteString("..") sb.WriteString(strconv.Itoa(b)) sb.WriteString(", ") } } // resize changes the length of S.data to n, keeping old values. // It preserves the invariant S.data[i] = 0, n ≤ i < cap(data). func (S *Set) resize(n int) { d := S.data if S.realloc(n) { low := min(S.min>>shift, len(d)) // Words below S.min>>shift are zero; copy(S.data[low:], d[low:]) // there is no need to copy them. } } // realloc creates a slice S.data of length n, possibly zeroing out old values. // It preserves the invariant S.data[i] = 0, n ≤ i < cap(data). // It returns true if new memory has been allocated. func (S *Set) realloc(n int) (didAlloc bool) { if c := cap(S.data); c < n { S.data = make([]uint64, n, newCap(n, c)) return true } // Add zeroes if shrinking. d := S.data low := max(n, S.min>>shift) // The words d[i], i < low, are already zero. for i := len(d) - 1; i >= low; i-- { d[i] = 0 } S.data = d[:n] return false } // newCap suggests a new increased capacity, favoring powers of two, // when growing a slice to length n. The suggested capacities guarantee // linear amortized cost for repeated memory allocations. func newCap(n, prevCap int) int { return max(n, nextPow2(prevCap)) } // nextPow2 returns the smallest p = 1, 2, 4, ..., 2^k such that p > n, // or MaxInt if p > MaxInt. func nextPow2(n int) (p int) { if n <= 0 { return 1 } if k := MaxPos(uint64(n)) + 1; k < BitsPerWord-1 { return 1 << uint(k) } return MaxInt } // trim slices S.data by removing all trailing words equal to zero. // It preserves the invariant: S.min = min(S) if S ≠ ∅, 0 otherwise. func (S *Set) trim() { d := S.data i := len(d) - 1 for i >= 0 && d[i] == 0 { i-- } S.data = d[:i+1] if i == -1 { S.min = 0 } } // findMinFrom finds the mininum element starting the search from n. // It assumes that such an element exists; however, it returns 0 if // len(data) == 0. func findMinFrom(n int, data []uint64) int { if len(data) == 0 { return 0 } i := n >> shift w := data[i] for w == 0 { i++ w = data[i] } return i<> uint(bpw-1-(n-m)) << uint(m) } func min(a, b int) int { if a < b { return a } return b } func max(a, b int) int { if a > b { return a } return b }