More Operations on Sorted Sets

C++

Java

C++

Java

In Java, iterators are helpful for looping through sets.

Iterators used with HashSet would yield the elements in random order:

Copy
Set<Integer> set = new HashSet<Integer>();
set.add(1);
set.add(3);
set.add(0);
set.add(-2);
Iterator it = set.iterator();
while (it.hasNext()) {
	Integer i = (Integer)it.next();
	System.out.print(i + " ");  // returns some random order
}

But with TreeSet the elements are in sorted order:

Copy
Set<Integer> set = new TreeSet<Integer>();
set.add(1);
set.add(3);
set.add(0);
set.add(-2);
Iterator it = set.iterator();
while (it.hasNext()) {
	Integer i = (Integer)it.next();
	System.out.print(i + " ");  // returns -2 0 1 3
}

Instead of creating an iterator and looping with it like in C++, Java provides a for-each loop which creates a hidden iterator and loops with it automatically:

Copy
Set<Integer> set = new TreeSet<Integer>();
set.add(1);
set.add(3);
set.add(0);
set.add(-2);
for (int i : set) {
	System.out.print(i + " ");  // returns -2 0 1 3
}

Python

In Python, we can use iter() to obtain the iterator object of any iterable. Using next() on the iterator lets you iterate through the iterable.

Copy
nums = {1, "three", 0, (1, 3, 5), 4}
iterator = iter(nums)

print(next(iterator))  # 0
print(next(iterator))  # 1
print(next(iterator))  # 4
print(next(iterator))  # three
print(next(iterator))  # (1, 3, 5)

# Note:
# The numbers appear ordered when iterating due to the hashing implementation
# Sets are not ordered, so do not expect this behavior

Python's iterators are fundamental to iteration and are used in for loops and tuple unpacking, for example. This is useful when you want lower level control over your iteration.

In competitive programming however, you will find that the itertools library provides more common uses that can save you lots of time. The most useful of these are shown below.

Copy
import itertools

# Iterate through and accumulate with a specified function
nums = [1, 0, 5, 2, -4, 3]
psa = list(itertools.accumulate(nums))  # operator.add is the default
print(psa)  # [1, 1, 6, 8, 4, 7]
print(list(itertools.accumulate(nums, min)))  # [1, 0, 0, 0, -4, -4]
print(list(itertools.accumulate(nums, max)))  # [1, 1, 5, 5, 5, 5]

# Iterate through the perms and combs below

C++

The sorted std::set also supports:

• lower_bound: returns an iterator to the least element greater than or equal to some element k.
• upper_bound: returns an iterator to the least element strictly greater than some element k.

Copy
set<int> s;
s.insert(1);                        // [1]
s.insert(14);                       // [1, 14]
s.insert(9);                        // [1, 9, 14]
s.insert(2);                        // [1, 2, 9, 14]
cout << *s.upper_bound(7) << '\n';  // 9
cout << *s.upper_bound(9) << '\n';  // 14
cout << *s.lower_bound(5) << '\n';  // 9
cout << *s.lower_bound(9) << '\n';  // 9
cout << *s.begin() << '\n';         // 1
auto it = s.end();
cout << *(--it) << '\n';    // 14
s.erase(s.upper_bound(6));  // [1, 2, 14]

Java

TreeSets in Java allow for a multitude of additional operations:

• first(): returns the lowest element in the set
• last(): returns the greatest element in the set
• lower(E v): returns the greatest element strictly less than v
• floor(E v): returns the greatest element less than or equal to v
• higher(E v): returns the least element strictly greater than v
• ceiling(E v): returns the least element greater than or equal to v

Copy
TreeSet<Integer> set = new TreeSet<Integer>();
set.add(1);                          // [1]
set.add(14);                         // [1, 14]
set.add(9);                          // [1, 9, 14]
set.add(2);                          // [1, 2, 9, 14]
System.out.println(set.higher(7));   // 9
System.out.println(set.higher(9));   // 14
System.out.println(set.lower(5));    // 2
System.out.println(set.first());     // 1
System.out.println(set.last());      // 14
set.remove(set.higher(6));           // [1, 2, 14]
System.out.println(set.higher(23));  // ERROR, no such element exists

Python

Python does not have a sorted set, so see the C++ or Java if you want an implementation from the standard library.

Creating a sorted set from scratch is time consuming, and not reccomended to do so in contests, but if you are curious regarding the implementation, you can find an AVL tree implementation of a sorted set below.

Since some online judges such as LeetCode include additional libraries, there is also an implementation of sorted sets included from the sortedcontainers library. All operations shown below are $\mathcal{O}(\log N)$ time, except for its $\mathcal{O}(N \log N)$ initialization.

Copy
from sortedcontainers import SortedSet

sorted_set = SortedSet([5, 1, 3, 2])
print(sorted_set)  # SortedSet([1, 2, 3, 4, 7])

# Add elements
sorted_set.add(4)
sorted_set.add(6)
print(sorted_set)  # SortedSet([1, 2, 3, 4, 5, 6])

C++

The ordered map also allows:

• lower_bound: returns the iterator pointing to the lowest entry not less than the specified key
• upper_bound: returns the iterator pointing to the lowest entry strictly greater than the specified key respectively.

Copy
map<int, int> m;
m[3] = 5;     // [(3, 5)]
m[11] = 4;    // [(3, 5); (11, 4)]
m[10] = 491;  // [(3, 5); (10, 491); (11, 4)]
cout << m.lower_bound(10)->first << " " << m.lower_bound(10)->second << '\n';  // 10 491
cout << m.upper_bound(10)->first << " " << m.upper_bound(10)->second << '\n';  // 11 4
m.erase(11);  // [(3, 5); (10, 491)]
if (m.upper_bound(10) == m.end()) {
	cout << "end" << endl;  // Prints end
}

Java

The ordered map additionally supports firstKey / firstEntry and lastKey / lastEntry, returning the lowest key/entry and the highest key/entry, as well as higherKey / higherEntry and lowerKey / lowerEntry, returning the lowest key/entry strictly higher than the specified key, or the highest key/entry strictly lower than the specified key.

Copy
TreeMap<Integer, Integer> map = new TreeMap<Integer, Integer>();
map.put(3, 5);                           // [(3, 5)]
map.put(11, 4);                          // [(3, 5); (11, 4)]
map.put(10, 491);                        // [(3, 5); (10, 491); (11, 4)]
System.out.println(map.firstKey());      // 3
System.out.println(map.firstEntry());    // (3, 5)
System.out.println(map.lastEntry());     // (11, 4)
System.out.println(map.higherEntry(4));  // (10, 491)
map.remove(11);                          // [(3, 5); (10, 491)]
System.out.println(map.lowerKey(4));     // 3
System.out.println(map.lowerKey(3));     // ERROR

Python

Python does not have built in sorted maps or sorted sets in it's standard library. However, sorted maps in Python can be created by adding a dictionary to a sorted set, where each element in the sorted set is a key in the dictionary and values can be assigned with the dictionary.

This is the most straightforward implementation, and the ground up implementation of a sorted set can be found in the above section.

Since some online judges include additional libraries, an implementation of SortedDict from the sortedcontainers library can be found below. All operations shown below are $\mathcal{O}(\log N)$ time, except for its $\mathcal{O}(N \log N)$ initialization and getting all items, keys, or values which is $\mathcal{O}(N)$ time.

Copy
from sortedcontainers import SortedDict

sorted_map = SortedDict({1: "one", 4: "four", 3: "three"})
print(sorted_map)  # SortedDict({1: 'one', 3: 'three', 4: 'four'})

# Add elements
sorted_map[2] = "two"
sorted_map[5] = "five"

# Output SortedDict({1: 'one', 2: 'two', 3: 'three', 4: 'four', 5: 'five'})

C++

In addition to all of the regular set operations,

• the count() method returns the number of times an element is present in the multiset. However, this method takes time linear in the number of matches so you shouldn't use it in a contest.
• To remove a value once, use ms.erase(ms.find(val)).
• To remove all occurrences of a value, use ms.erase(val).

Copy
multiset<int> ms;
ms.insert(1);                  // [1]
ms.insert(14);                 // [1, 14]
ms.insert(9);                  // [1, 9, 14]
ms.insert(2);                  // [1, 2, 9, 14]
ms.insert(9);                  // [1, 2, 9, 9, 14]
ms.insert(9);                  // [1, 2, 9, 9, 9, 14]
cout << ms.count(4) << '\n';   // 0
cout << ms.count(9) << '\n';   // 3
cout << ms.count(14) << '\n';  // 1
ms.erase(ms.find(9));
cout << ms.count(9) << '\n';  // 2
ms.erase(9);
cout << ms.count(9) << '\n';  // 0

Java

While there is no multiset in Java, we can implement one using the TreeMap from values to their respective frequencies. We declare the TreeMap implementation globally so that we can write functions for adding and removing elements from it. The first, last, higher, and lower operations still function as intended; just use firstKey, lastKey, higherKey, and lowerKey respectively.

Copy
static TreeMap<Integer, Integer> multiset = new TreeMap<Integer, Integer>();

public static void main(String[] args) { ... }

static void add(int x) {
	if (multiset.containsKey(x)) {
		multiset.put(x, multiset.get(x) + 1);
	} else {
		multiset.put(x, 1);
	}

C++

C++

Copy
priority_queue<int> pq;
pq.push(7);                // [7]
pq.push(2);                // [2, 7]
pq.push(1);                // [1, 2, 7]
pq.push(5);                // [1, 2, 5, 7]
cout << pq.top() << endl;  // 7
pq.pop();                  // [1, 2, 5]
pq.pop();                  // [1, 2]
pq.push(6);                // [1, 2, 6]

Java

Java

In Java (unlike in C++), we delete and retrieve the lowest element.

Copy
PriorityQueue<Integer> pq = new PriorityQueue<Integer>();
pq.add(7);                      // [7]
pq.add(2);                      // [7, 2]
pq.add(1);                      // [7, 2, 1]
pq.add(5);                      // [7, 5, 2, 1]
System.out.println(pq.peek());  // 1
pq.poll();                      // [7, 5, 2]
pq.poll();                      // [7, 5]
pq.add(6);                      // [7, 6, 5]

Python

Python

In Python (unlike in C++), we delete and retrieve the lowest element.

Note that Python's priority queue is not encapsulated; heapq operates on a provided list directly by turning it into a heap, then doing operations on the heap.

Copy
import heapq

pq = []

"""
The following line is not necessary since pq starts as an empty list. However,
for lists of length greater than 1, heapify is required to turn the list into a
heap.
"""
heapq.heapify(pq)

C++

Copy
#include <bits/stdc++.h>
using namespace std;

#define sz(x) (x).size()

string s;
int m;
set<int> dif;
multiset<int> ret;

Java

Copy
import java.io.*;
import java.util.*;

class BitInversion {
	public static TreeMap<Integer, Integer> ret = new TreeMap<Integer, Integer>();
	public static String s;
	public static int m;
	public static Set<Integer> dif = new TreeSet<Integer>();

	public static void modify(int x) {

C++

Copy
#include <bits/stdc++.h>
using namespace std;

#define sz(x) (int)(x).size()

string s;
int m;
set<int> dif;
priority_queue<int> ret;
int cnt[200005];

Java

Copy
import java.io.*;
import java.util.*;

class BitInversion {
	public static PriorityQueue<Integer> pq =
	    new PriorityQueue<Integer>(Collections.reverseOrder());
	public static String s;
	public static int m;
	public static TreeSet<Integer> dif = new TreeSet<Integer>();
	public static int cnt[];