Python中的枚举和2D数组？

Question

I'm currently working on some code in Python, and I have the following enumerations:

#Enumeration - matter type
matterType = ['matter', 'antimatter']
enumMatterType = enumerate(matterType)

#Enumeration  - flavor
flavor = ['electron', 'mu', 'tau']
enumFlavor = enumerate(flavor)

I've also created a 2D array:

alpha = [[alpha_electron, alpha_anti_electron], 
        [alpha_mu, alpha_anti_mu], 
        [alpha_tau, alpha_anti_tau]]

Which is supposed to be an array of constants that I have defined earlier. I'm using these constants in an equation, but I'm referencing them as alpha[0][0] or alpha[0][1] , etc. I feel like when someone looks at my code, they would understand it better if I had something like alpha[matter][electron] or alpha[antimatter][electron] . Is there any way to use the enumeration from earlier to define my variables from the array like this? How can I use enumeration with arrays in general? How would I be able to write my alpha values with the enumeration? Please let me know if any part of this does not make sense - I would be happy to clarify.

Answer 1

You should consider a dictionary with tuple keys:

matterType = ['matter', 'antimatter']
flavor = ['electron', 'mu', 'tau']

alpha = [['alpha_electron', 'alpha_anti_electron'], 
        ['alpha_mu', 'alpha_anti_mu'], 
        ['alpha_tau', 'alpha_anti_tau']]

ref_dict = {}
for x, pair_item in enumerate(flavor):
    for y, item in enumerate(matterType):
        ref_dict[(item, pair_item)] = alpha[x][y]

Now let's say I want "matter" and "mu":

In[60]: ref_dict[('matter', 'mu')]
Out[60]: 'alpha_mu'

Answer 2

If you want to create a set of enumeration constants, that's exactly what the enum module is for.

In your case, you want enumeration constants that have an integer value, so you can use them as indexes into a list. For that, you want IntEnum :

from enum import IntEnum

class MatterType(IntEnum):
    MATTER = 0
    ANTIMATTER = 1

class Flavor(IntEnum):
    ELECTRON = 0
    MU = 1
    TAU = 2

And now, instead of this:

alpha[0][0]

… you can write this:

alpha[MatterType.MATTER][Flavor.ELECTRON]