at last 339C

Author: kumar80

339C.cpp

@@ -50,44 +50,64 @@
                                              `''''`'''i==_+=_=i__
                                                      ||'''- '    `.
                                                       `-.......-''
-*/                              
-#include<bits/stdc++.h>
+*/
+#include <bits/stdc++.h>
 using namespace std;
- 
-#define fastio ios_base::sync_with_stdio(0); cin.tie(0)
-#define LL long long 
-#define mod 998244353 
-#define FOR(i, j, k) for (int i=j ; i<k ; i++)
-#define ROF(i, j, k) for (int i=j ; i>=k ; i--) 
- 
-const long long INF = 1e18;
-const long long MAX  = 1e5+10;
-
-int main(){
-    fastio;
-    int t=1; //cin>>t;
-    while(t--){
-    	int m; string s; cin>>s>>m;
-    	std::vector<int> v; FOR(i,0,s.size()) if(s[i]=='1') v.push_back(i+1);
-     		LL prev = 0;
-    std::vector<int> res;
-    FOR(j,0,v.size()){
-      std::vector<int> ans; ans.push_back(v[j]);    LL l=v[j],r=0;   prev=v[j]; //cout<<j<<"; ";
-    	FOR(i,1,m){
-    		int diff = abs(l-r);
-    		auto it = upper_bound(v.begin(), v.end(),diff);
-    		if(it!=v.end() && *it==prev) it = upper_bound(v.begin(), v.end(),prev);
 
-    		if(it==v.end()) break ;//{cout<<"NO"; return 0;}
-    		if(i&1) r+=*it; 
-    		else l+=*it;
-    		ans.push_back(*it); prev = *it;
-    	} 	
-      if(ans.size()==m) {res=ans;break;}
-    }
-    if(res.size()!=m) {cout<<"NO\n"; return 0;}
+#define fastio                  \
+  ios_base::sync_with_stdio(0); \
+  cin.tie(0)
+#define LL long long
+#define mod 998244353
+#define FOR(i, j, k) for (auto i = j; i < k; i++)
+#define ROF(i, j, k) for (auto i = j; i >= k; i--)
 
-    	cout<<"YES\n";
-    	for(auto i : res) std::cout<<i<<" "; 
+const long long INF = 1e18;
+const long long MAX = 1e5 + 10;
+bool pos = false;
+int m;
+void f(vector<int> &v, int x, int j, vector<int> p)
+{
+  if (j == m && !pos)
+  {
+    pos = true;
+    cout << "YES\n";
+    for (auto i : p)
+      cout << i << " ";
+  }
+  if (pos || j > m)
+    return;
+  auto it = upper_bound(v.begin(), v.end(), abs(x));
+  FOR(i, it, v.end())
+  {
+    vector<int> temp = p;
+    if (temp.size() == 0)
+    {
+      temp.push_back(*i), f(v, x + *i, j + 1, temp);
+      continue;
     }
+    if ((j & 1) && (*i != *temp.rbegin()))
+      temp.push_back(*i), f(v, x - *i, j + 1, temp);
+    else if ((*i != *temp.rbegin()) && (j % 2 == 0))
+      temp.push_back(*i), f(v, x + *i, j + 1, temp);
+  }
+   
+}
+int main()
+{
+  fastio;
+  int t = 1; //cin>>t;
+  while (t--)
+  {
+    string s;
+    cin >> s >> m;
+    std::vector<int> v;
+    FOR(i, 0, s.size())
+    if (s[i] == '1')
+      v.push_back(i + 1);
+    vector<int> p;
+    f(v, 0, 0, p);
+    if (!pos)
+      cout << "NO";
+  }
 }