1% Functions that operate on sets in the form of arrays and lists:
2% Copyright (C) 2010-2017,2018 John E. Davis
3%
4% This file is part of the S-Lang Library and may be distributed under the
5% terms of the GNU General Public License.  See the file COPYING for
6% more information.
7%
8%  Functions: unique, union, complement, intersection, ismember
9private define pop_set_object ()
10{
11   variable a = ();
12   if ((typeof(a) != Array_Type) && (typeof(a) != List_Type))
13     a = [a];
14   return a;
15}
16
17private define list_unique (a)
18{
19   variable len = length(a);
20   variable indices = Int_Type[len];
21   variable i, j, k;
22
23   k = 0;
24   _for i (0, len-1, 1)
25     {
26	variable a_i = a[i];
27	_for j (0, i-1, 1)
28	  {
29	     if (_eqs(a_i, a[j]))
30	       break;
31	  }
32	then
33	  {
34	     indices[k] = i;
35	     k++;
36	  }
37     }
38   return indices[[0:k-1]];
39}
40
41
42define unique ()
43{
44   variable i, j, len;
45   variable a;
46
47   if (_NARGS != 1)
48     {
49	_pop_n (_NARGS);
50	usage ("i = unique (a); %% i = indices of unique elements of a");
51     }
52
53   a = pop_set_object ();
54   if (typeof(a) == List_Type)
55     {
56	try
57	  {
58	     a = list_to_array (a);
59	  }
60	catch AnyError: return list_unique (a);
61     }
62
63   len = length(a);
64   if (len <= 1)
65     return [0:len-1];
66
67   if (length(array_shape(a)) != 1)
68     a = _reshape (__tmp(a),[len]);
69
70   try
71     {
72	i = array_sort(a);
73     }
74   catch AnyError: return list_unique (a);
75
76   a = a[i];
77   if (a[0] == a[-1])		       %  all equal
78     return [0];
79   j = where (shift(a,-1)!=a);
80   % Now, i contains the sorted indices, and j contains the indices into the
81   % sorted array.  So, the unique elements are given by a[i][j] where a is
82   % the original input array.  It seems amusing that the indices given by
83   % [i][j] are also given by i[j].
84   return i[__tmp(j)];
85}
86
87define union ()
88{
89   !if (_NARGS)
90     usage ("U = union (A, B, ..., C);");
91
92   variable args = {}, obj;
93   variable has_list = 0;
94   loop (_NARGS)
95     {
96	has_list += (typeof (obj) == List_Type);
97	obj = pop_set_object ();
98	list_insert (args, obj);
99     }
100
101   variable a = NULL;
102   if (has_list == 0)
103     {
104	try
105	  {
106	     a = [__push_list (args)];
107	  }
108	catch AnyError:;
109     }
110
111   if (a == NULL)
112     {
113	a = {};
114	foreach obj (args)
115	  {
116	     if (typeof(obj) == List_Type)
117	       {
118		  list_join (a, obj);
119		  continue;
120	       }
121	     foreach (obj)
122	       {
123		  variable x = ();
124		  list_append (a, x);
125	       }
126	  }
127     }
128   return a[unique (a)];
129}
130
131% return indices of a that are not in b
132private define list_complement (a, b)
133{
134   variable lena = length(a), lenb = length(b);
135   variable indices = Int_Type[lena];
136   variable i, j, k;
137
138   k = 0;
139   _for i (0, lena-1, 1)
140     {
141	variable a_i = a[i];
142	_for j (0, lenb-1, 1)
143	  {
144	     if (_eqs(a_i, b[j]))
145	       break;
146	  }
147	then
148	  {
149	     indices[k] = i;
150	     k++;
151	  }
152     }
153   return indices[[0:k-1]];
154}
155
156define complement ()
157{
158   variable a, b;
159   if (_NARGS != 2)
160     usage ("\
161i = complement (a, b);\n\
162%% Returns the indices of the elements of `a' that are not in `b'");
163
164   b = pop_set_object ();
165   a = pop_set_object ();
166
167   variable
168     lena = length(a),
169     lenb = length(b);
170
171   if ((lena == 0) || (lenb == 0))
172     return [0:lena-1];
173
174   variable sia, sib;
175   try
176     {
177	if (typeof (a) == List_Type)
178	  a = list_to_array (a);
179	if (typeof (b) == List_Type)
180	  b = list_to_array (b);
181	sia = array_sort (a);
182	sib = array_sort (b);
183     }
184   catch AnyError:
185     return list_complement (a, b);
186
187   variable
188     c = Int_Type [lena], j = 0,
189     ia, ib, xa, xb, k;
190
191   ia = 0; ib = 0;
192   xb = b[sib[ib]];
193   while (ia < lena)
194     {
195	k = sia[ia];
196	xa = a[k];
197	if (xa < xb)
198	  {
199	     c[j] = k;
200	     j++;
201	     ia++;
202	     continue;
203	  }
204	if (xb == xa)
205	  {
206	     ia++;
207	     continue;
208	  }
209
210	while (ib++, (ib < lenb) && (xa > b[sib[ib]]))
211	  ;
212	if (ib == lenb)
213	  {
214	     variable n = lena-ia;
215	     c[[j:j+n-1]] = sia[[ia:lena-1]];
216	     j += n;
217	     break;
218	  }
219	xb = b[sib[ib]];
220	if (xa == xb)
221	  ia++;
222     }
223   return c[[0:j-1]];
224}
225
226% Return the indices into a of the common elements of both a and b
227define intersection ()
228{
229   if (_NARGS < 2)
230     usage ("\
231i = intersection (a, b, .., c);\n\
232%% Returns the indices of 'a' of the common elements of b,.., c");
233
234   variable b = pop_set_object ();
235   loop (_NARGS-1)
236     {
237	variable a = pop_set_object ();
238	variable i = complement (a, __tmp(b));
239	i = complement (a, a[i]);
240	b = a[i];
241     }
242   return i;
243}
244
245% Returns whether or not a is a member of b.
246define ismember ()
247{
248   if (_NARGS != 2)
249     usage ("I = ismember (a, b);\n\
250Returns a boolean array indicated whether the corresponding elements of 'a'\n\
251are members of 'b'");
252
253   variable a, b;
254   (a, b) = ();
255   if ((typeof(a) == Array_Type) || (typeof(a) == List_Type))
256     {
257	variable lena = length (a);
258	variable result = Char_Type[lena];
259	result[intersection(a,b)] = 1;
260	return result;
261     }
262   return 0 != length (intersection (a, b));
263}
264