python - Dataframe 获得某行的最大值对应的的列索引（包含多个相等的最大值的情况）

import pandas as pd 
import numpy as np

# create a dataframe
table = pd.DataFrame(np.zeros([3,2]), index=["s1", "s2", "s3"], columns=["a1","a2"])
"""
     a1   a2
s1  0.0  0.0
s2  0.0  0.0
s3  0.0  0.0
"""
# 获得 s1 行
s1 = table.loc["s1",:]

# 获得 s1 中最大值的索引，可能有多个
s1_argmax = s1[s1 == s1.max()].index

# randomly choose 1 index
s1_argmax = np.random.choice(s1_argmax)