Fixup NHot function

pyeda/boolalg/expr.py

@@ -509,12 +509,14 @@ def NHot(n, *xs, simplify=True):
         fstr = "expected 0 <= n <= {}, got {}"
         raise ValueError(fstr.format(len(xs), n))
 
-    xs = {Expression.box(x).node for x in xs}
+    xs = [Expression.box(x).node for x in xs]
+    num = len(xs)
     terms = list()
-    for hots in itertools.combinations(xs, n):
-        colds = xs - set(hots)
-        not_colds = tuple(map(exprnode.not_, colds))
-        terms.append(exprnode.and_(*(hots + not_colds)))
+    for hot_idxs in itertools.combinations(range(num), n):
+        hot_idxs = set(hot_idxs)
+        _xs = [xs[i] if i in hot_idxs else exprnode.not_(xs[i])
+               for i in range(num)]
+        terms.append(exprnode.and_(*_xs))
     y = exprnode.or_(*terms)
     if simplify:
         y = y.simplify()

pyeda/boolalg/test/test_expr.py

@@ -11,7 +11,7 @@
     Expression,
     Zero, One,
     Not, Or, And, Nor, Nand, Xor, Xnor, Equal, Unequal, Implies, ITE,
-    OneHot0, OneHot, Majority, AchillesHeel, Mux,
+    OneHot0, OneHot, NHot, Majority, AchillesHeel, Mux,
 )
 
 
@@ -174,6 +174,16 @@ def test_onehot():
     assert OneHot(a, b, c, conj=False).equivalent((~a | ~b) & (~a | ~c) & (~b | ~c) & (a | b | c))
     assert OneHot(a, b, c, conj=True).equivalent((~a | ~b) & (~a | ~c) & (~b | ~c) & (a | b | c))
 
+def test_nhot():
+    assert NHot(2, 0, 0, 0) is Zero
+    assert NHot(2, 0, 0, 1) is Zero
+    assert NHot(2, 0, 1, 0) is Zero
+    assert NHot(2, 0, 1, 1) is One
+    assert NHot(2, 1, 0, 0) is Zero
+    assert NHot(2, 1, 0, 1) is One
+    assert NHot(2, 1, 1, 0) is One
+    assert NHot(2, 1, 1, 1) is Zero
+
 def test_majority():
     assert Majority(0, 0, 0) is Zero
     assert Majority(0, 0, 1) is Zero