Testing compress with various iterables

from itertools import compress from collections import OrderedDict, deque from operator import itemgetter import numpy as np import heapq # Boolean # r = compress(False, [0,0,1,0,1,1]) # TypeError: 'bool' object is not iterable # Char # Only the first value in the mask is considered r = compress('X', [0,0,1,0,1,1]) for el in r: print(el) # Empty result r = compress('X', [1,0,1,0,1,1]) for el in r: print(el) # X # String r = compress('ABCDEF', [0,0,1,0,1,1]) for el in r: print(el) # C # E # F print(itemgetter(2,4,5)('ABCDEF')) # ('C', 'E', 'F') # List r = compress([10, 20, 30,40, 50, 60], [0,0,1,0,1,1]) for el in r: print(el) # 30 # 50 # 60 # Tuple r = compress((('Andrew' , 19), ('Boris', 22), ('Christina', 20), ('Denise', 22), ('Eric', 19), ('Daniel', 20)), [0,0,1,0,1,1]) for el in r: print(el) # ('Christina', 20) # ('Eric', 19) # ('Daniel', 20) # Set r = compress({'apple', 'banana', 'orange', 'strawberry', 'avocado', 'lemon'}, [0,0,1,0,1,1]) for el in r: print(el) # Doesn't work as expected # avocado # orange # banana # Dictionary r = compress({'a': 1, 'b': 2, 'c': 3, 'd': 4, 'e': 5, 'f': 6}, [0,0,1,0,1,1]) for el in r: print(el) # Doesn't work as expected # b # d # f # Ordered dictionary d = OrderedDict() d['a'] = 1 d['b'] = 2 d['c'] = 3 d['d'] = 4 d['e'] = 5 d['f'] = 6 r = compress(d, [0,0,1,0,1,1]) for el in r: print(el) # Now it works as expected # c # e # f # Double-ended queue (deque) d = deque('abcdef') r = compress(d, [0,0,1,0,1,1]) for el in r: print(el) # c # e # f for _ in range(3): # Test compress after rotations d.rotate(1) r = compress(d, [0,0,1,0,1,1]) for el in r: print(el) # 1. 2. 3. # b a f # d c b # e d c # Heap queue h = [] for item in [5,4,8,6,3,4]: heapq.heappush(h, item) # Now h is [3, 4, 4, 6, 5, 8] r = compress(h, [0,0,1,0,1,1]) for el in r: print(el) # 4 # 5 # 8 # Numpy array arr = np.array(range(1,7)) r = compress(arr, [0,0,1,0,1,1]) for el in r: print(el) # 3 # 5 # 6 print(np.take(arr, [2,4,5])) # [3 5 6] # Numpy 2D array attempt arr = np.array(range(1,17)).reshape(4,4) # Generate mask with 16 zeros or ones at random mask = np.random.randint(2, size=16).reshape(4,4) r = compress(arr, mask) #print([el for el in r]) # fails for 2D rows = cols = range(4) print(np.array([arr[i,j] if mask[i,j] else 0 for i in rows for j in cols]).reshape(4,4)) ''' [[ 0 2 0 0] [ 5 6 7 0] [ 0 10 11 0] [ 0 0 0 16]] ''' # Back to 1D, using compress arr = np.array(range(1,17)) mask = np.random.randint(2, size=16) r = compress(arr, mask) print(np.array([el for el in r])) # Example: [ 3 5 6 8 9 10 11 13 14 16] # We don't know the number of elements in advance, so we can't automatically reshape to a 4x4 matrix at the end # Numpy masked array import numpy.ma as ma print(ma.masked_array(arr, mask).reshape(4,4)) ''' Example: [[-- 2 -- --] [5 6 -- --] [9 10 11 12] [-- 14 15 --]] '''