I’m struggling to simplify my irregular nested np.where clauses. Is there a way to make the code more readable?
JavaScript
x
26
26
1
df["COL"] = np.where(
2
(df["A1"] == df["B1"]) & (df["A1"].notna()),
3
np.where(
4
(df["A1"] == df["C"]),
5
np.where(
6
(df["A"] == df["B"]) & df["A"].notna() & (df["A"] != df["A1"]),
7
"Text1",
8
df["A1"]
9
),
10
"Text2"
11
),
12
np.where(
13
(df["A"] == df["B"]) & (df["A"].notna()),
14
np.where(
15
(df["A"] == df["C"]),
16
df["A"],
17
"Text1"
18
),
19
np.where(
20
(df["C"].notna()),
21
df["C"],
22
"Text3"
23
)
24
)
25
)
26
Advertisement
Answer
Using np.select
as suggested by @sammywemmy:
JavaScript
1
22
22
1
# Create boolean masks
2
m1 = (df["A1"] == df["B1"]) & (df["A1"].notna())
3
m11 = (df["A1"] == df["C"])
4
m12 = (df["A"] == df["B"]) & (df["A"].notna())
5
m111 = (df["A"] == df["B"]) & df["A"].notna() & (df["A"] != df["A1"])
6
m121 = (df["A"] == df["C"])
7
m122 = (df["C"].notna())
8
9
# Combine them
10
condlist = [m1 & m11 & m111,
11
m1 & m11 & ~m111,
12
m1 & ~m11,
13
~m1 & m12 & m121,
14
~m1 & m12 & ~m121
15
~m1 & ~m12 & m122,
16
~m1 & ~m12 & ~m122]
17
18
# Values for each combination
19
choicelist = ["Text1", df["A1"], "Text2", df["A"], "Text1", df["C"], "Text3"]
20
21
out = np.select(condlist, choicelist)
22