CODEWITHSUNDEEP
python
max
max

Reputation: 4521

Better way to test membership in multiple sets?

My code here is terrible, is there a pythonic way to improve it?

I have three sets:

set1 = {1,2,3}
set2 = {2,3,4}
set3 = {4,5,6}

and a combined set

combined = set1 | set2 | set3

I need a dict that returns whether a number is each set.

For example

d[2] = {'set1':True, 'set2':True, 'set3':False}

My code is like this:

def in_set(num):
    d = {}
    if num in set1:
        d['set1'] = True
    else:
        d['set1'] = False

    if num in set2:
        d['set2'] = True
    else:
        d['set2'] = False

    if num in set3:
        d['set3'] = True
    else:
        d['set3'] = False

    return d

Upvotes: 3

Views: 323

Answers (5)

martineau
martineau

Reputation: 123501

Here's something flexible and user-friendly utilizing functools.reduce():

from functools import reduce

# Update and return dictionary d.
update_dict = lambda d, key, value: (d.update({key: value}), d)[1]

def in_sets(num, **kwargs):
    return reduce(lambda d, kv: update_dict(d, kv[0], num in kv[1]), kwargs.items(), {})

if __name__ == '__main__':

    test1 = in_sets(3, set1={1,2,3}, set2={2,3,4}, set3={4,5,6})
    print(test1)  # -> {'set1': True, 'set2': True, 'set3': False}

    test2 = in_sets(5, set1={1,2,3}, set2={2,3,4}, set3={4,5,6})
    print(test2)  # -> {'set1': False, 'set2': False, 'set3': True}

Upvotes: 0

ekiim
ekiim

Reputation: 842

When it comes to the pythonic way of doing things, depends on a slightly broader picture,

For instance, you can create a function, that given an element a collection of sets, it will return which sets contains it, something like this.

def contained_in_which(elem, list_of_sets):
    temp = []
    for i, s in enumerate(list_of_sets):
        if elem in s:
            temp.append(i)

This could be one way, not very pythonic but clean.

If you want to get more sophisticated you could do something like this.

Depending on the case, this seems a little bit more generic, to use, because It's just an extra attribute for the set, and you can do stuff like this, The interesting part is that arranging a little better your data structure you can extend a lot of the defaults.

class NamedSet(set):
    def __init__(self, *args, name="SET"):
        if isinstance(args[0], set) and len(args) == 1:
            args = list(args[0])    
        self.name = name
        super().__init__(args)

set1 = NamedSet({1,2,3}, name="ONE")
set2 = NamedSet({2,3,4}, name="TWO")
set3 = NamedSet({4,5,6}, name="THREE")
setList = [set1, set2, set3]

n = 2 
for i in filter(lambda i_set: 2 in i_set, setList):
    print(i.name)

Upvotes: 1

LeKhan9
LeKhan9

Reputation: 1350

If you make a list of all your sets, then you can do something like this:

sets = [set1, set2, set3]

num = 1
print {'set{}'.format(i+1): num in set_i for i, set_i in enumerate(sets)}

The comprehension isn't the cleanest, but its not too bad either :)

output:

{'set1': True, 'set2': False, 'set3': False}

Upvotes: 1

Tasos Papastylianou
Tasos Papastylianou

Reputation: 22255

Start from a dictionary in the first place; naming individual variables in a var1, var2, var3 pattern usually indicates you need some kind of container in the first place anyway.

Then it becomes a simple dict comprehension:

Sets = {'set1' : {1,2,3}, 'set2' : {2,3,4}, 'set3' : {3,4,5} };
def in_set( x ): return {i : x in Sets[i] for i in Sets }

in_set( 2 )  #>  {'set2': True, 'set1': True, 'set3': False}

Upvotes: 2

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 477309

You can simplify this with:

def in_set(num):
    return {
        'set1': num in set1,
        'set2': num in set2,
        'set3': num in set3
    }

You can even make a dictionary of the set names that map on these sets and then use dictionary comprehension:

my_sets = {
    'set1': set1,
    'set2': set2,
    'set3': set3
}

def in_set(num):
    return {
        k: num in v
        for k, v in my_sets.items()
    }

Here we can easily change the sets (or their names), and thus for example work with more/other sets.

Upvotes: 6

Related Questions